From 783f2c60b2a6c5fcd566a511aebf94691e310609 Mon Sep 17 00:00:00 2001 From: Duc Ngo Date: Wed, 8 Aug 2018 14:39:05 -0700 Subject: [PATCH 1/5] nomnigraph - Enhancements to subgraph matching APIs (#10218) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/10218 SubtreeMatchCriteria now supports: - nonTerminal flag : if this is set, it means we only match the root of the subtree and do not care about the children. Example use case: to match an "input" node but does not care how the input is produced. Additional tests for these new logic are added to subgraph_matcher_test.cc. Subgraph matching APIs for NNGraph is also added. (Further enhancement to make the SubgraphMatching API constructs a Subgraph object/more diagnostic information will go later). Reviewed By: bwasti Differential Revision: D9156092 fbshipit-source-id: 3f28ac15d9edd474b3e0cd51fd7e6f973299d061 --- .../nomnigraph/Representations/NeuralNet.cc | 18 ++++ .../include/nomnigraph/Graph/Graph.h | 4 + .../nomnigraph/Representations/NeuralNet.h | 53 +++++++++++ .../include/nomnigraph/Support/Common.h | 1 + .../Transformations/SubgraphMatcher.h | 53 ++++++----- .../core/nomnigraph/tests/neural_net_test.cc | 92 +++++++++++++++++++ .../nomnigraph/tests/subgraph_matcher_test.cc | 75 ++++++++------- 7 files changed, 240 insertions(+), 56 deletions(-) create mode 100644 caffe2/core/nomnigraph/tests/neural_net_test.cc diff --git a/caffe2/core/nomnigraph/Representations/NeuralNet.cc b/caffe2/core/nomnigraph/Representations/NeuralNet.cc index 8281b1d3aa9df..8a54ed893d416 100644 --- a/caffe2/core/nomnigraph/Representations/NeuralNet.cc +++ b/caffe2/core/nomnigraph/Representations/NeuralNet.cc @@ -181,6 +181,24 @@ void coalesceInsertedDataDependencies(repr::NNModule* m) { } } +bool hasSingleOutputAndConsumer(NNGraph::NodeRef nodeRef) { + auto nodeOutputs = nn::getOutputs(nodeRef); + NOM_REQUIRE_OR_RET_FALSE(nodeOutputs.size() == 1); + auto nodeConsumers = nn::getConsumers(nodeOutputs.front()); + return nodeConsumers.size() == 1; +} + +NNNodeMatchCriteria matchAnyNode() { + return [](NNGraph::NodeRef /* unused */) { return true; }; +} + +NNSubtree operatorTree( + const NNNodeMatchCriteria& root, + const std::vector& childrenCriteria, + int count) { + return NNSubtree(matchAnyNode(), {NNSubtree(root, childrenCriteria)}, count); +} + } // namespace nn } // namespace repr diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h b/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h index cb7b90059a156..25ee5e99a64ab 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Graph/Graph.h @@ -399,6 +399,10 @@ class Graph { return result; } + const size_t getNodesCount() const { + return (size_t)nodes_.size(); + } + const std::vector getMutableEdges() { std::vector result; for (auto& e : edges_) { diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h b/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h index 745a155b4a572..9f548b701d6b2 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Representations/NeuralNet.h @@ -17,6 +17,7 @@ #include "nomnigraph/Representations/ControlFlow.h" #include "nomnigraph/Support/Casting.h" #include "nomnigraph/Support/Pointer.h" +#include "nomnigraph/Transformations/SubgraphMatcher.h" #include #include @@ -420,6 +421,58 @@ void coalesceInsertedDataDependencies(repr::NNModule* m); template struct NodeHelper {}; +using NNNodeMatchCriteria = std::function; +using NNSubtree = nom::matcher::SubtreeMatchCriteria; + +bool hasSingleOutputAndConsumer(NNGraph::NodeRef nodeRef); + +template +NNNodeMatchCriteria matchNodeTypeWithPredicate( + const std::function predicate, + bool expectedSingleOutputAndConsumer = false) { + return + [&predicate, expectedSingleOutputAndConsumer](NNGraph::NodeRef nodeRef) { + NOM_REQUIRE_OR_RET_FALSE(is(nodeRef)); + if (expectedSingleOutputAndConsumer) { + NOM_REQUIRE_OR_RET_FALSE(hasSingleOutputAndConsumer(nodeRef)); + } + NodeType* node = get(nodeRef); + return predicate(nodeRef, *node); + }; +}; + +template +NNNodeMatchCriteria matchNodeType( + bool expectedSingleOutputAndConsumer = false) { + return [expectedSingleOutputAndConsumer](NNGraph::NodeRef nodeRef) { + if (expectedSingleOutputAndConsumer) { + NOM_REQUIRE_OR_RET_FALSE(hasSingleOutputAndConsumer(nodeRef)); + } + return is(nodeRef); + }; +} + +NNNodeMatchCriteria matchAnyNode(); + +struct NNNodeMatch { + static bool isMatch( + const NNGraph::NodeRef& node, + const NNNodeMatchCriteria& criteria) { + return criteria(node); + } +}; + +using NNSubgraphMatcher = + nom::matcher::SubgraphMatcher; + +// This helper method makes it easy to create matching criteria in NNGraph. +// For example, operatorTree(opMatch, ...) will refer to a tree like this: +// ... -> opMatch -> opMatch_Output +NNSubtree operatorTree( + const NNNodeMatchCriteria& root, + const std::vector& childrenCriteria = {}, + int count = 1); + } // namespace nn } // namespace repr diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Support/Common.h b/caffe2/core/nomnigraph/include/nomnigraph/Support/Common.h index 380afe4815dca..cef1bdec522a5 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Support/Common.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Support/Common.h @@ -31,6 +31,7 @@ #define NOM_REQUIRE_OR_BREAK(_cond) NOM_REQUIRE_OR_(_cond, break) #define NOM_REQUIRE_OR_RET_NULL(_cond) NOM_REQUIRE_OR_(_cond, return nullptr) #define NOM_REQUIRE_OR_RET_FALSE(_cond) NOM_REQUIRE_OR_(_cond, return false) +#define NOM_REQUIRE_OR_RET_TRUE(_cond) NOM_REQUIRE_OR_(_cond, return true) #define NOM_REQUIRE_OR_RET(_cond) NOM_REQUIRE_OR_(_cond, return ) // Implements accessors for a generic type T. If the type is not diff --git a/caffe2/core/nomnigraph/include/nomnigraph/Transformations/SubgraphMatcher.h b/caffe2/core/nomnigraph/include/nomnigraph/Transformations/SubgraphMatcher.h index 08ead74295074..b02ef9210bcf9 100644 --- a/caffe2/core/nomnigraph/include/nomnigraph/Transformations/SubgraphMatcher.h +++ b/caffe2/core/nomnigraph/include/nomnigraph/Transformations/SubgraphMatcher.h @@ -1,6 +1,9 @@ #ifndef NOM_TRANFORMATIONS_SUBGRAPH_MATCHER_H #define NOM_TRANFORMATIONS_SUBGRAPH_MATCHER_H +#include +#include + namespace nom { namespace matcher { @@ -10,8 +13,10 @@ namespace matcher { * - Node matching criteria for the subtree's root. * - Children subtree matching criteria * - A count, which means we may want more than one of this subtree. The count - * can be unlimited. The count is only used when we match children of a - * subtree root, not matching the subtree itself. + * can be unlimited. The count is only used when we match children of a subtree + * root, not matching the subtree itself. + * - If nonTerminal flag is set, it means we only match the root and do not + * care about the children. */ template class SubtreeMatchCriteria { @@ -19,14 +24,26 @@ class SubtreeMatchCriteria { static const int kStarCount = -1; SubtreeMatchCriteria( const NodeMatchCriteria& root, - const std::vector& children, - int count) - : root_(root), children_(children), count_(count){}; + const std::vector& children = {}, + int count = 1, + bool nonTerminal = false) + : root_(root), + children_(children), + count_(count), + nonTerminal_(nonTerminal){}; + + // Non terminal + static SubtreeMatchCriteria nonTerminal( + const NodeMatchCriteria& root, + int count = 1) { + return SubtreeMatchCriteria(root, {}, count, true); + } private: NodeMatchCriteria root_; std::vector children_; int count_; + bool nonTerminal_; template friend class SubgraphMatcher; @@ -58,6 +75,11 @@ struct SubgraphMatcher { if (!isNodeMatch(root, criteria.root_)) { return false; } + if (criteria.nonTerminal_) { + // This is sufficient to be a match if this criteria specifies a non + // terminal node. + return true; + } auto& edges = invertGraphTraversal ? root->getInEdges() : root->getOutEdges(); @@ -87,9 +109,9 @@ struct SubgraphMatcher { (isStarCount || countMatch < expectedCount); currentEdgeIdx++) { auto edge = edges[currentEdgeIdx]; - auto nextNode = invertGraphTraversal ? edge->tail() : edge->head(); + auto child = invertGraphTraversal ? edge->tail() : edge->head(); - if (!isSubtreeMatch(nextNode, childrenCriteria, invertGraphTraversal)) { + if (!isSubtreeMatch(child, childrenCriteria, invertGraphTraversal)) { if (!isStarCount) { // If the current criteria isn't a * pattern, this indicates a // failure. @@ -150,23 +172,6 @@ struct SubgraphMatcher { } }; -// Convenient methods to create subtree matching criteria. -template -SubtreeMatchCriteria tree( - const NodeMatchCriteria& root, - const std::vector>& children = {}, - int count = 1) { - return SubtreeMatchCriteria(root, children, count); -} - -template -SubtreeMatchCriteria treeStar( - const NodeMatchCriteria& root, - const std::vector>& children = {}) { - return tree( - root, children, SubtreeMatchCriteria::kStarCount); -} - } // namespace matcher } // namespace nom diff --git a/caffe2/core/nomnigraph/tests/neural_net_test.cc b/caffe2/core/nomnigraph/tests/neural_net_test.cc new file mode 100644 index 0000000000000..a85b84a74b83c --- /dev/null +++ b/caffe2/core/nomnigraph/tests/neural_net_test.cc @@ -0,0 +1,92 @@ +#include + +#include "test_util.h" + +#include "nomnigraph/Representations/NeuralNet.h" +#include "nomnigraph/Support/Pointer.h" +#include "nomnigraph/Transformations/SubgraphMatcher.h" + +#include + +using namespace nom; +using namespace nom::repr; +using namespace nom::repr::nn; + +// Test for the NNGraph subgraph matching APIs. +TEST(NeuralNetGraph, ReplaceGraph) { + NNGraph graph; + + auto input1 = graph.createNode(util::make_unique("input1")); + auto input2 = graph.createNode(util::make_unique("input2")); + auto sum = graph.createNode(util::make_unique()); + auto sumOutput = graph.createNode(util::make_unique("sumOutput")); + auto relu = graph.createNode(util::make_unique()); + auto reluOutput = graph.createNode(util::make_unique("reluOutput")); + + graph.createEdge(input1, sum); + graph.createEdge(input2, sum); + graph.createEdge(sum, sumOutput); + graph.createEdge(sumOutput, relu); + graph.createEdge(relu, reluOutput); + + /* input1 input2 + \ / + \ / + sum + | + | + sumOutput + | + relu + | + reluOutput + */ + + // clang-format off + auto pattern = NNSubtree( + matchNodeType(), { + operatorTree( + matchNodeType(), { + NNSubtree::nonTerminal(matchNodeType(), 2) + }), + }); + // clang-format on + + EXPECT_FALSE(NNSubgraphMatcher::isSubtreeMatch(sum, pattern)); + EXPECT_FALSE(NNSubgraphMatcher::isSubtreeMatch(reluOutput, pattern)); + EXPECT_FALSE(NNSubgraphMatcher::isSubtreeMatch(input1, pattern)); + + EXPECT_TRUE(NNSubgraphMatcher::isSubtreeMatch(relu, pattern)); + + NNSubgraphMatcher::replaceSubtree( + graph, pattern, [](NNGraph& g, NNGraph::NodeRef relu) { + auto sumOutput = getInputs(relu)[0]; + auto sum = getProducer(sumOutput); + + auto fusedNode = g.createNode(util::make_unique()); + g.deleteNode(sumOutput); + g.replaceNode(relu, fusedNode); + g.replaceNode(sum, fusedNode); + + g.deleteNode(sum); + g.deleteNode(relu); + + return true; + }); + + /* + Fused graph: + + input1 input2 + \ / + \ / + sumRelu + | + | + output + */ + EXPECT_EQ(graph.getNodesCount(), 4); + auto fusedNode = getProducer(reluOutput); + EXPECT_TRUE(is(fusedNode)); + EXPECT_EQ(getInputs(fusedNode).size(), 2); +} diff --git a/caffe2/core/nomnigraph/tests/subgraph_matcher_test.cc b/caffe2/core/nomnigraph/tests/subgraph_matcher_test.cc index ddd8a15fcdc2b..3441ee445a3ce 100644 --- a/caffe2/core/nomnigraph/tests/subgraph_matcher_test.cc +++ b/caffe2/core/nomnigraph/tests/subgraph_matcher_test.cc @@ -25,6 +25,7 @@ struct TestNodeMatch { using TestGraph = Graph; using TestMatcher = SubgraphMatcher; +using Tree = SubtreeMatchCriteria; Criteria any() { return Criteria("*"); @@ -37,7 +38,7 @@ SubtreeMatchCriteria operatorTree( const Criteria& root, const std::vector>& childrenCriteria = {}, int count = 1) { - return tree(any(), {tree(root, childrenCriteria)}, count); + return Tree(any(), {Tree(root, childrenCriteria)}, count); } std::map TestGraphNodePrinter( @@ -160,18 +161,18 @@ struct DataFlowTestGraph { SubtreeMatchCriteria DataFlowTestGraphCriteria() { // clang-format off - return tree( + return Tree( Criteria("opG"),{ operatorTree("opF", { // Note: we currently don't enforce that these 2 opC nodes // have to be the same. operatorTree("opB", { operatorTree("opC", { - treeStar(Criteria("input")) + Tree(Criteria("input"), {}, Tree::kStarCount) }, 2), }) }), - tree(any()) // matches dataI + Tree(any()) // matches dataI }); // clang-format on } @@ -223,20 +224,30 @@ TEST(SubgraphMatcher, IsSubtreeMatch) { N3 N4 N6 N7 */ - auto subtree = tree(any(), {tree(any()), tree(any())}); + auto subtree = Tree(any(), {Tree(any()), Tree(any())}); EXPECT_FALSE(TestMatcher::isSubtreeMatch(n1, subtree, false)); EXPECT_FALSE(TestMatcher::isSubtreeMatch(n4, subtree, false)); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n2, subtree, false)); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n5, subtree, false)); - subtree = tree(Criteria("5"), {tree(any()), tree(any())}); + subtree = Tree(Criteria("5"), {Tree(any()), Tree(any())}); EXPECT_FALSE(TestMatcher::isSubtreeMatch(n2, subtree, false)); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n5, subtree, false)); - subtree = tree(any(), {tree(any()), tree(Criteria("4"))}); + subtree = Tree(any(), {Tree(any()), Tree(Criteria("4"))}); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n2, subtree, false)); EXPECT_FALSE(TestMatcher::isSubtreeMatch(n5, subtree, false)); + + // Accepts non terminal node + subtree = Tree(any(), {Tree::nonTerminal(any()), Tree::nonTerminal(any())}); + EXPECT_TRUE(TestMatcher::isSubtreeMatch(n1, subtree, false)); + EXPECT_TRUE(TestMatcher::isSubtreeMatch(n2, subtree, false)); + EXPECT_TRUE(TestMatcher::isSubtreeMatch(n5, subtree, false)); + EXPECT_FALSE(TestMatcher::isSubtreeMatch(n3, subtree, false)); + EXPECT_FALSE(TestMatcher::isSubtreeMatch(n4, subtree, false)); + EXPECT_FALSE(TestMatcher::isSubtreeMatch(n6, subtree, false)); + EXPECT_FALSE(TestMatcher::isSubtreeMatch(n7, subtree, false)); } // Test subtree matching in which * (repeated) matching of children is allowed. @@ -259,49 +270,49 @@ TEST(SubgraphMatcher, IsSubtreeMatchRepeated) { graph.createEdge(n1, n5B); graph.createEdge(n1, n5C); - auto subtree = tree(any(), {tree(Criteria("2"))}); + auto subtree = Tree(any(), {Tree(Criteria("2"))}); EXPECT_FALSE(TestMatcher::isSubtreeMatch(n1, subtree, false)); - subtree = tree(any(), {treeStar(Criteria("2"))}); + subtree = Tree(any(), {Tree(Criteria("2"), {}, Tree::kStarCount)}); EXPECT_FALSE(TestMatcher::isSubtreeMatch(n1, subtree, false)); // clang-format off - subtree = tree(any(), { - tree(Criteria("2")), - tree(Criteria("3"), {}, 2), - tree(Criteria("4"), {}, 2), - tree(Criteria("5"), {}, 3) + subtree = Tree(any(), { + Tree(Criteria("2")), + Tree(Criteria("3"), {}, 2), + Tree(Criteria("4"), {}, 2), + Tree(Criteria("5"), {}, 3) }); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n1, subtree, false)); - subtree = tree(any(), { - tree(Criteria("2")), - tree(Criteria("3"), {}, 2), - tree(Criteria("4"), {}, 2), - treeStar(Criteria("5")) + subtree = Tree(any(), { + Tree(Criteria("2")), + Tree(Criteria("3"), {}, 2), + Tree(Criteria("4"), {}, 2), + Tree(Criteria("5"), {}, Tree::kStarCount) }); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n1, subtree, false)); - subtree = tree(any(), { - tree(Criteria("2")), - treeStar(Criteria("3")), - tree(Criteria("4"), {}, 2), - treeStar(Criteria("5")) + subtree = Tree(any(), { + Tree(Criteria("2")), + Tree(Criteria("3"), {}, Tree::kStarCount), + Tree(Criteria("4"), {}, 2), + Tree(Criteria("5"), {}, Tree::kStarCount) }); EXPECT_TRUE(TestMatcher::isSubtreeMatch(n1, subtree, false)); - subtree = tree(any(), { - tree(Criteria("2")), - treeStar(Criteria("3")), + subtree = Tree(any(), { + Tree(Criteria("2")), + Tree(Criteria("3"), {}, Tree::kStarCount), }); // Fails because there are unmatched edges. EXPECT_FALSE(TestMatcher::isSubtreeMatch(n1, subtree, false)); - subtree = tree(any(), { - tree(Criteria("2")), - tree(Criteria("3"), {}, 2), - tree(Criteria("4")), - tree(Criteria("5"), {}, 3) + subtree = Tree(any(), { + Tree(Criteria("2")), + Tree(Criteria("3"), {}, 2), + Tree(Criteria("4")), + Tree(Criteria("5"), {}, 3) }); // Fails because the count is wrong; we have 2 edges to node N4 while // the pattern expects only 1. From 6e49f933ad01ccd9064e5ef3d433c7aefcfe3d86 Mon Sep 17 00:00:00 2001 From: Thomas Viehmann Date: Wed, 8 Aug 2018 21:04:44 -0700 Subject: [PATCH 2/5] Check that result is on CPU for CPU unary ops kernels (#10358) Summary: Fixes: #10270 Pull Request resolved: https://github.com/pytorch/pytorch/pull/10358 Differential Revision: D9233066 Pulled By: soumith fbshipit-source-id: 39b7524fe55ddb899fb27e2c0ef504ce54dbad35 --- aten/src/ATen/native/cpu/UnaryOpsKernel.cpp | 1 + test/test_torch.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp index 459838a9b6c68..7ecedff060bf2 100644 --- a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp +++ b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp @@ -104,6 +104,7 @@ static void sigmoid_kernel(Tensor& result, const Tensor& self) { #define IMPLEMENT_FLOAT_KERNEL(dispatchtypes, op) \ static void op##_kernel(Tensor& result, const Tensor& self) { \ + checkBackend(#op, {result}, kCPU); \ AT_DISPATCH_##dispatchtypes##_TYPES(self.type(), #op, [&] { \ if (self.is_contiguous() && result.is_contiguous()) { \ vml::v##op( \ diff --git a/test/test_torch.py b/test/test_torch.py index e494981abf031..a9bff95e0bccf 100644 --- a/test/test_torch.py +++ b/test/test_torch.py @@ -597,6 +597,12 @@ def test_floor(self): def test_ceil(self): self._test_math_by_name('ceil') + @unittest.skipIf(not torch.cuda.is_available(), 'no CUDA') + def test_ceil_out_cpu_cuda(self): + a = torch.randn(1) + b = torch.randn(1, device="cuda") + self.assertRaises(RuntimeError, lambda: torch.ceil(a, out=b)) + def test_rsqrt(self): def rsqrt(x): if x == 0: From 9dfc4edc68c2d6ba2b98008b50e48b28e3e50c25 Mon Sep 17 00:00:00 2001 From: Marat Dukhan Date: Wed, 8 Aug 2018 22:00:18 -0700 Subject: [PATCH 3/5] Update NNPACK and cpuinfo submodules (#8564) Summary: Bring in extra optimizations in Winograd-based convolution on NEON Pull Request resolved: https://github.com/pytorch/pytorch/pull/8564 Reviewed By: hlu1 Differential Revision: D9088140 Pulled By: Maratyszcza fbshipit-source-id: 2089191416db98bdad8f0e4848b1435fcf74a88b --- third_party/NNPACK | 2 +- third_party/cpuinfo | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/NNPACK b/third_party/NNPACK index 3eb0d453662d0..af40ea7d12702 160000 --- a/third_party/NNPACK +++ b/third_party/NNPACK @@ -1 +1 @@ -Subproject commit 3eb0d453662d05a708f43b108bed9e17b705383e +Subproject commit af40ea7d12702f8ae55aeb13701c09cad09334c3 diff --git a/third_party/cpuinfo b/third_party/cpuinfo index 1e6c8c99d27f2..4e8f04355892c 160000 --- a/third_party/cpuinfo +++ b/third_party/cpuinfo @@ -1 +1 @@ -Subproject commit 1e6c8c99d27f2b5eb9d2e6231055c6a4115b85e5 +Subproject commit 4e8f04355892c5deb64a51731a6afdb544a4294d From 037d8d1bab9b952aecf084b1a6b6e70be4fefe34 Mon Sep 17 00:00:00 2001 From: Tongzhou Wang Date: Wed, 8 Aug 2018 22:20:14 -0700 Subject: [PATCH 4/5] Order Loss functions alphabetically in nn.rst Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/10365 Differential Revision: D9237287 Pulled By: SsnL fbshipit-source-id: 28e9de76b9cfd8f63c8df561ff1531ea8d0803ea --- docs/source/nn.rst | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/nn.rst b/docs/source/nn.rst index 9b3563f83b2d4..68420d837bf80 100644 --- a/docs/source/nn.rst +++ b/docs/source/nn.rst @@ -1182,6 +1182,11 @@ Loss functions .. autofunction:: binary_cross_entropy +:hidden:`binary_cross_entropy_with_logits` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autofunction:: binary_cross_entropy_with_logits + :hidden:`poisson_nll_loss` ~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1247,11 +1252,6 @@ Loss functions .. autofunction:: nll_loss -:hidden:`binary_cross_entropy_with_logits` -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. autofunction:: binary_cross_entropy_with_logits - :hidden:`smooth_l1_loss` ~~~~~~~~~~~~~~~~~~~~~~~~ From 04f381650e8cabe55305bd04706446f2af296816 Mon Sep 17 00:00:00 2001 From: Tongzhou Wang Date: Wed, 8 Aug 2018 23:55:36 -0700 Subject: [PATCH 5/5] Resubmit: Fix dataloader hang when it is not completely iterated (#10366) Summary: https://github.com/pytorch/pytorch/pull/9655 Pull Request resolved: https://github.com/pytorch/pytorch/pull/10366 Differential Revision: D9237393 Pulled By: SsnL fbshipit-source-id: fabfad7f371ba33300098f6b885c0e3f26c3e14a --- test/test_dataloader.py | 116 +++++++++++++++++++-------------- torch/utils/data/dataloader.py | 85 ++++++++++++------------ 2 files changed, 109 insertions(+), 92 deletions(-) diff --git a/test/test_dataloader.py b/test/test_dataloader.py index bb61cced71753..3c4fe7540ac03 100644 --- a/test/test_dataloader.py +++ b/test/test_dataloader.py @@ -205,9 +205,12 @@ class SleepDataset(Dataset): def __init__(self, size, sleep_sec): self.size = size self.sleep_sec = sleep_sec + self.sleeped = False def __getitem__(self, idx): - time.sleep(self.sleep_sec) + if not self.sleeped: + time.sleep(self.sleep_sec) + self.sleeped = True return idx def __len__(self): @@ -251,7 +254,7 @@ def __len__(self): def _test_timeout(): - dataset = SleepDataset(10, 10) + dataset = SleepDataset(10, 3) dataloader = DataLoader(dataset, batch_size=2, num_workers=2, timeout=1) _ = next(iter(dataloader)) @@ -478,36 +481,40 @@ def test_error_workers(self): @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") @unittest.skipIf(TEST_WITH_ROCM, "test doesn't currently work on the ROCm stack") def test_partial_workers(self): - "check that workers exit even if the iterator is not exhausted" - loader = iter(DataLoader(self.dataset, batch_size=2, num_workers=4, pin_memory=True)) - workers = loader.workers - worker_manager_thread = loader.worker_manager_thread - for i, sample in enumerate(loader): - if i == 3: - break - del loader - for w in workers: - w.join(JOIN_TIMEOUT) - self.assertFalse(w.is_alive(), 'subprocess not terminated') - self.assertEqual(w.exitcode, 0) - worker_manager_thread.join(JOIN_TIMEOUT) - self.assertFalse(worker_manager_thread.is_alive()) + r"""Check that workers exit even if the iterator is not exhausted.""" + for pin_memory in (True, False): + loader = iter(DataLoader(self.dataset, batch_size=2, num_workers=4, pin_memory=pin_memory)) + workers = loader.workers + if pin_memory: + pin_memory_thread = loader.pin_memory_thread + for i, sample in enumerate(loader): + if i == 10: + break + del loader + for w in workers: + w.join(JOIN_TIMEOUT) + self.assertFalse(w.is_alive(), 'subprocess not terminated') + if pin_memory: + pin_memory_thread.join(JOIN_TIMEOUT) + self.assertFalse(pin_memory_thread.is_alive()) @staticmethod - def _manager_process(dataset, worker_pids, manager_exit_event): + def _main_process(dataset, worker_pids, main_exit_event, raise_error): loader = iter(DataLoader(dataset, batch_size=2, num_workers=4, pin_memory=True)) workers = loader.workers for i in range(len(workers)): worker_pids[i] = int(workers[i].pid) for i, sample in enumerate(loader): if i == 3: - break - # Simulate a dirty exit of the manager process - manager_exit_event.set() - if IS_WINDOWS: - os.system('taskkill /PID ' + str(os.getpid()) + ' /F') - else: - os.kill(os.getpid(), signal.SIGKILL) + # Simulate an exit of the manager process + main_exit_event.set() + if raise_error: + raise RuntimeError('Error') + else: + if IS_WINDOWS: + os.system('taskkill /PID ' + str(os.getpid()) + ' /F') + else: + os.kill(os.getpid(), signal.SIGKILL) @staticmethod def _is_process_alive(pid, pname): @@ -532,32 +539,43 @@ def _is_process_alive(pid, pname): @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") @unittest.skipIf(TEST_WITH_ROCM, "test doesn't currently work on the ROCm stack") def test_main_process_unclean_exit(self): - '''There might be ConnectionResetError or leaked semaphore warning (due to dirty process exit), \ + r'''There might be ConnectionResetError or leaked semaphore warning (due to dirty process exit), \ but they are all safe to ignore''' - worker_pids = mp.Array('i', [0] * 4) - manager_exit_event = mp.Event() - p = mp.Process(target=TestDataLoader._manager_process, - args=(self.dataset, worker_pids, manager_exit_event)) - p.start() - - manager_exit_event.wait() - - exit_status = [False] * len(worker_pids) - start_time = time.time() - pname = 'python' - while True: - for i in range(len(worker_pids)): - pid = worker_pids[i] - if not exit_status[i]: - if not TestDataLoader._is_process_alive(pid, pname): - exit_status[i] = True - if all(exit_status): - break - else: - time.sleep(1) - self.assertFalse(time.time() - start_time > MANAGER_STATUS_CHECK_INTERVAL + JOIN_TIMEOUT, - 'subprocess not terminated') + # `raise_error` controls if the main process is KILL-ed by OS or just + # simply raises an error. Both cases are interesting because + # 1. In case of it is KILL-ed by OS, the workers need to automatically + # discover that their parent is dead and exit gracefully. + # 2. In case of it raises an error itself, the parent process needs to + # take care of exiting the worker and then exits itself gracefully. + for raise_error in (True, False): + worker_pids = mp.Array('i', [0] * 4) + + main_exit_event = mp.Event() + p = mp.Process(target=TestDataLoader._main_process, + args=(self.dataset, worker_pids, main_exit_event, raise_error)) + p.start() + worker_pids[-1] = p.pid + + main_exit_event.wait() + + exit_status = [False] * len(worker_pids) + start_time = time.time() + pname = 'python' + while True: + for i in range(len(worker_pids)): + pid = worker_pids[i] + if not exit_status[i]: + if not TestDataLoader._is_process_alive(pid, pname): + exit_status[i] = True + if all(exit_status): + break + else: + if time.time() - start_time > MANAGER_STATUS_CHECK_INTERVAL + JOIN_TIMEOUT: + self.fail('subprocess not terminated') + time.sleep(1) + p.join(MANAGER_STATUS_CHECK_INTERVAL + JOIN_TIMEOUT - (time.time() - start_time)) + self.assertFalse(p.is_alive(), 'main process not terminated') def test_len(self): def check_len(dl, expected): @@ -601,7 +619,7 @@ def __len__(self): self.assertIsInstance(batch, tt) @unittest.skipIf(not TEST_NUMPY, "numpy unavailable") - def test_default_colate_bad_numpy_types(self): + def test_default_collate_bad_numpy_types(self): import numpy as np # Should be a no-op diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py index 10457a6653602..60789e9fb6299 100644 --- a/torch/utils/data/dataloader.py +++ b/torch/utils/data/dataloader.py @@ -72,7 +72,7 @@ def is_alive(self): return os.getppid() == self.manager_pid -def _worker_loop(dataset, index_queue, data_queue, collate_fn, seed, init_fn, worker_id): +def _worker_loop(dataset, index_queue, data_queue, done_event, collate_fn, seed, init_fn, worker_id): global _use_shared_memory _use_shared_memory = True @@ -86,6 +86,11 @@ def _worker_loop(dataset, index_queue, data_queue, collate_fn, seed, init_fn, wo random.seed(seed) torch.manual_seed(seed) + # Do not wait for putting thread to join when this worker exits. Otherwise, + # this worker may always be waiting to put and doesn't check index_queue + # and done_event for termination signal. + data_queue.cancel_join_thread() + if init_fn is not None: init_fn(worker_id) @@ -95,11 +100,13 @@ def _worker_loop(dataset, index_queue, data_queue, collate_fn, seed, init_fn, wo try: r = index_queue.get(timeout=MANAGER_STATUS_CHECK_INTERVAL) except queue.Empty: - if watchdog.is_alive(): + if watchdog.is_alive() and not done_event.is_set(): continue else: break - if r is None: + # use done_event so that we can get faster exiting signal even if there + # are still indices in index_queue + if r is None or done_event.is_set(): break idx, batch_indices = r try: @@ -111,7 +118,7 @@ def _worker_loop(dataset, index_queue, data_queue, collate_fn, seed, init_fn, wo del samples -def _worker_manager_loop(in_queue, out_queue, done_event, pin_memory, device_id): +def _pin_memory_loop(in_queue, out_queue, done_event, pin_memory, device_id): if pin_memory: torch.cuda.set_device(device_id) @@ -122,7 +129,7 @@ def _worker_manager_loop(in_queue, out_queue, done_event, pin_memory, device_id) if done_event.is_set(): return raise - if r is None: + if r is None or done_event.is_set(): break if isinstance(r[1], ExceptionWrapper): out_queue.put(r) @@ -242,7 +249,6 @@ def __init__(self, loader): self.num_workers = loader.num_workers self.pin_memory = loader.pin_memory and torch.cuda.is_available() self.timeout = loader.timeout - self.done_event = threading.Event() self.sample_iter = iter(self.batch_sampler) @@ -252,35 +258,32 @@ def __init__(self, loader): self.worker_init_fn = loader.worker_init_fn self.index_queues = [multiprocessing.Queue() for _ in range(self.num_workers)] self.worker_queue_idx = 0 - self.worker_result_queue = multiprocessing.SimpleQueue() + self.worker_result_queue = multiprocessing.Queue() self.batches_outstanding = 0 self.worker_pids_set = False self.shutdown = False self.send_idx = 0 self.rcvd_idx = 0 self.reorder_dict = {} + self.done_event = multiprocessing.Event() self.workers = [ multiprocessing.Process( target=_worker_loop, args=(self.dataset, self.index_queues[i], - self.worker_result_queue, self.collate_fn, base_seed + i, + self.worker_result_queue, self.done_event, + self.collate_fn, base_seed + i, self.worker_init_fn, i)) for i in range(self.num_workers)] - if self.pin_memory or self.timeout > 0: + if self.pin_memory: self.data_queue = queue.Queue() - if self.pin_memory: - maybe_device_id = torch.cuda.current_device() - else: - # do not initialize cuda context if not necessary - maybe_device_id = None - self.worker_manager_thread = threading.Thread( - target=_worker_manager_loop, + self.pin_memory_thread = threading.Thread( + target=_pin_memory_loop, args=(self.worker_result_queue, self.data_queue, self.done_event, self.pin_memory, - maybe_device_id)) - self.worker_manager_thread.daemon = True - self.worker_manager_thread.start() + torch.cuda.current_device())) + self.pin_memory_thread.daemon = True + self.pin_memory_thread.start() else: self.data_queue = self.worker_result_queue @@ -366,33 +369,29 @@ def __getstate__(self): raise NotImplementedError("_DataLoaderIter cannot be pickled") def _shutdown_workers(self): - try: - if not self.shutdown: - self.shutdown = True - self.done_event.set() - for q in self.index_queues: - q.put(None) - # if some workers are waiting to put, make place for them - try: - while not self.worker_result_queue.empty(): - self.worker_result_queue.get() - except (FileNotFoundError, ImportError): - # Many weird errors can happen here due to Python - # shutting down. These are more like obscure Python bugs. - # FileNotFoundError can happen when we rebuild the fd - # fetched from the queue but the socket is already closed - # from the worker side. - # ImportError can happen when the unpickler loads the - # resource from `get`. - pass - # done_event should be sufficient to exit worker_manager_thread, - # but be safe here and put another None - self.worker_result_queue.put(None) - finally: - # removes pids no matter what + if not self.shutdown: + self.shutdown = True + # removes pids from the C side data structure first so worker + # termination afterwards won't trigger false positive error report. if self.worker_pids_set: _remove_worker_pids(id(self)) self.worker_pids_set = False + self.done_event.set() + if self.pin_memory: + # Sending `None` to `pin_memory_thread` must be before + # stopping worker processes because the workers may leave + # corrupted data in `worker_result_queue`, causing + # `pin_memory_thread` unable to read and terminate properly. + self.worker_result_queue.put(None) + # Workers can't be waiting to put be cause their output queue + # is a multiprocessing.Queue and its .put is non-blocking. + # They can only be waiting to get, so we put `None` here. + for q in self.index_queues: + q.put(None) + for w in self.workers: + w.join() + if self.pin_memory: + self.pin_memory_thread.join() def __del__(self): if self.num_workers > 0: