From fd5eb3b351780e34893c7c3ccf2319db1846e58d Mon Sep 17 00:00:00 2001 From: Doris Xin Date: Fri, 17 May 2019 15:05:36 -0700 Subject: [PATCH] Remove dropout nodes during inference (#1018) * initial version. there's a bug * allow nodes with multiple outputs to be removed if only one output is used by downstream ops * move node output checks into their own methods * add test data file * address comments by @kkaranasos and @skottmckay * address more comments by @kkaranasos * a comment to clarify the position of the mask output --- onnxruntime/core/graph/graph_utils.cc | 32 ++++++++++-- onnxruntime/core/graph/graph_utils.h | 8 ++- .../core/optimizer/dropout_elimination.cc | 43 +++++++++++++++ .../core/optimizer/dropout_elimination.h | 31 +++++++++++ onnxruntime/test/ir/utils_test.cc | 49 ++++++++++++++++++ .../test/optimizer/graph_transform_test.cc | 25 +++++++++ .../test/testdata/transform/dropout.onnx | Bin 0 -> 718 bytes 7 files changed, 183 insertions(+), 5 deletions(-) create mode 100644 onnxruntime/core/optimizer/dropout_elimination.cc create mode 100644 onnxruntime/core/optimizer/dropout_elimination.h create mode 100644 onnxruntime/test/testdata/transform/dropout.onnx diff --git a/onnxruntime/core/graph/graph_utils.cc b/onnxruntime/core/graph/graph_utils.cc index 6d45b65f3d..d3a21293a5 100644 --- a/onnxruntime/core/graph/graph_utils.cc +++ b/onnxruntime/core/graph/graph_utils.cc @@ -348,10 +348,36 @@ const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const Node& node, const s return iter == attrs.end() ? nullptr : &iter->second; } +/** Checks for nodes with >= 1 outputs, if only one of the outputs is input to downstream Operators. */ +static bool IsOnlyOneOutputUsed(const Node& node) { + if (node.GetOutputEdgesCount() > 1) { + const int unassigned = -1; + int first_output = unassigned; + for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) { + if (first_output == unassigned) { + first_output = it->GetSrcArgIndex(); + } else if (first_output != it->GetSrcArgIndex()) { + return false; + } + } + } + return true; +} + +bool IsOutputUsed(const Node& node, int index) { + for (auto it = node.OutputEdgesBegin(), end = node.OutputEdgesEnd(); it != end; ++it) { + if (it->GetSrcArgIndex() == index) { + return true; + } + } + return false; +} + bool RemoveNode(Graph& graph, Node& node) { - // Cannot remove a node with implicit inputs, with multiple output NodeArgs (multiple output edges is fine), - // or whose output is also a graph output. - if (!node.ImplicitInputDefs().empty() || node.OutputDefs().size() != 1 || graph.IsNodeOutputsInGraphOutputs(node)) { + // Cannot remove a node with implicit inputs, whose output is also a graph output, + // or with more than one of its outputs as input to downstream Operators. + if (!node.ImplicitInputDefs().empty() || + graph.IsNodeOutputsInGraphOutputs(node) || !IsOnlyOneOutputUsed(node)) { return false; } diff --git a/onnxruntime/core/graph/graph_utils.h b/onnxruntime/core/graph/graph_utils.h index cace5b9286..0f84a7575c 100644 --- a/onnxruntime/core/graph/graph_utils.h +++ b/onnxruntime/core/graph/graph_utils.h @@ -32,6 +32,9 @@ bool IsSupportedProvider(const Node& node, fed to multiple downstream operators, i.e., it can have multiple output edges. */ bool IsSingleInSingleOutNode(const Node& node); +/** Checks if the output at the specified index is input to downstream Nodes. */ +bool IsOutputUsed(const Node& node, int index); + /** Returns true if the graph has the given input.*/ bool IsGraphInput(const Graph& graph, const NodeArg* input); @@ -64,8 +67,9 @@ Status ForAllMutableSubgraphs(Graph& main_graph, std::function f Status ForAllSubgraphs(const Graph& main_graph, std::function func); /** Removes the given Node from the Graph and keeps Graph consistent by rebuilding needed connections. - We support the removal of the Node if it has no implicit inputs and a single output (but it can have multiple - output edges). As for the Node's inputs, we support the following cases: + We support the removal of the Node as long as the following conditions hold: + - There should be no implicit inputs. + - Only one of the outputs is used by downstream operators (but it can have multiple output edges). - If the Node has a single incoming node (and possibly multiple initializers), we can remove the Node and connect its incoming node to its outgoing nodes. - If the Node has a single initializer as input, we remove the Node and feed the initializer as input to its diff --git a/onnxruntime/core/optimizer/dropout_elimination.cc b/onnxruntime/core/optimizer/dropout_elimination.cc new file mode 100644 index 0000000000..f08643f029 --- /dev/null +++ b/onnxruntime/core/optimizer/dropout_elimination.cc @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/logging/logging.h" +#include "core/graph/graph_viewer.h" +#include "core/graph/op.h" +#include "core/graph/graph_utils.h" +#include "core/optimizer/rewrite_rule.h" +#include "core/optimizer/dropout_elimination.h" + +namespace onnxruntime { + +Status EliminateDropout::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect) { + if (graph_utils::RemoveNode(graph, node)) { + rule_effect = RewriteRuleEffect::kRemovedCurrentNode; + } + + return Status::OK(); +} + +bool EliminateDropout::SatisfyCondition(const Graph& graph, const Node& node) { + // We currently support elimination for Dropout operator v1, v6, v7, and v10. + if (!graph_utils::IsSupportedOptypeVersionAndDomain(node, "Dropout", {1, 6, 7, 10})) { + return false; + } + + if (graph.IsNodeOutputsInGraphOutputs(node)) { + return false; + } + + // A Dropout Node has one required output and an optional second output (i.e. at index = 1), `mask`. + // It can be safely removed if a) it has only one output + // or b) if the `mask` output is present but is not an input to any downstream Nodes. + // The `is_test` attribute in v1 and v6 is captured by the check for the `mask` output. + if (graph_utils::IsSingleInSingleOutNode(node)) { + return true; + } else { + // The `mask` output, which must not be used downstream for node removal, is at outputs[1]. + return !graph_utils::IsOutputUsed(node, 1); + } +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/dropout_elimination.h b/onnxruntime/core/optimizer/dropout_elimination.h new file mode 100644 index 0000000000..2310eaa436 --- /dev/null +++ b/onnxruntime/core/optimizer/dropout_elimination.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/rewrite_rule.h" + +namespace onnxruntime { + +/** +@Class EliminateDropout + +Rewrite rule that eliminates dropout nodes without downstream dependencies. + +It is attempted to be triggered only on nodes with op type "Dropout". +*/ +class EliminateDropout : public RewriteRule { + public: + EliminateDropout() noexcept : RewriteRule("EliminateDropout") {} + + std::vector TargetOpTypes() const noexcept override { + return {"Dropout"}; + } + + private: + bool SatisfyCondition(const Graph& graph, const Node& node) override; + + Status Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect) override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/test/ir/utils_test.cc b/onnxruntime/test/ir/utils_test.cc index 4c36b3113e..96b87974ea 100644 --- a/onnxruntime/test/ir/utils_test.cc +++ b/onnxruntime/test/ir/utils_test.cc @@ -303,5 +303,54 @@ TEST(GraphUtils, TestMultiEdgeRemovalNodes) { ASSERT_TRUE(nodes[4]->InputDefs()[0]->Name() == "id_0_in"); } +TEST(GraphUtils, TestMultiOutputRemoveNode) { + + Model model("MultiOutputRemovalGraph"); + auto& graph = model.MainGraph(); + + TypeProto float_tensor; + float_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + TypeProto bool_tensor; + bool_tensor.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL); + bool_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + + auto& do_0_in = graph.GetOrCreateNodeArg("do_0_in", &float_tensor); + auto& do_0_out = graph.GetOrCreateNodeArg("do_0_out", &float_tensor); + auto& do_0_out1 = graph.GetOrCreateNodeArg("do_0_out1", &bool_tensor); + auto& id_1_out = graph.GetOrCreateNodeArg("id_1_out", &float_tensor); + auto& id_2_out = graph.GetOrCreateNodeArg("id_2_out", &bool_tensor); + + std::vector nodes; + nodes.push_back(&graph.AddNode("do_0", "Dropout", "Dropout node 0", {&do_0_in}, {&do_0_out, &do_0_out1})); + nodes.push_back(&graph.AddNode("id_1", "Identity", "Identity node 1", {&do_0_out}, {&id_1_out})); + nodes.push_back(&graph.AddNode("id_2", "Identity", "Identity node 2", {&do_0_out1}, {&id_2_out})); + + std::vector graph_outputs; + graph_outputs.push_back(&id_2_out); + graph.SetOutputs(graph_outputs); + + auto status = graph.Resolve(); + ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); + + ASSERT_EQ(graph.NumberOfNodes(), 3); + + // Check inputs/outputs of do_0, id_1, id_2 + ASSERT_EQ(nodes[0]->GetOutputEdgesCount(), 2); + ASSERT_EQ(nodes[1]->GetInputEdgesCount(), 1); + ASSERT_EQ(nodes[2]->GetInputEdgesCount(), 1); + + // Try to remove do_0, which should return false + // because both outputs are consumed by downstream Operators. + ASSERT_FALSE(graph_utils::RemoveNode(graph, *nodes[0])); + + // Try removing do_0 after removing id_2, which should return true + // because it now has exactly one output consumed by downstream Operators. + ASSERT_TRUE(graph_utils::RemoveNode(graph, *nodes[1])); + ASSERT_FALSE(graph_utils::IsOutputUsed(*nodes[0], 0)); + ASSERT_TRUE(graph_utils::RemoveNode(graph, *nodes[0])); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 1e2f902eb8..1d2facbbd7 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -7,6 +7,7 @@ #include "core/optimizer/graph_transformer.h" #include "core/optimizer/graph_transformer_mgr.h" #include "core/optimizer/identity_elimination.h" +#include "core/optimizer/dropout_elimination.h" #include "core/optimizer/slice_elimination.h" #include "core/optimizer/unsqueeze_elimination.h" #include "core/optimizer/conv_bn_fusion.h" @@ -62,6 +63,30 @@ TEST(GraphTransformationTests, IdentityElimination) { ASSERT_TRUE(op_to_count["Identity"] == 0); } +TEST(GraphTransformationTests, DropoutEliminationSingleOutput) { + string model_uri = MODEL_FOLDER + "dropout.onnx"; + std::shared_ptr model; + ASSERT_TRUE(Model::Load(model_uri, model).IsOK()); + Graph& graph = model->MainGraph(); + std::map op_to_count = CountOpsInGraph(graph); + ASSERT_TRUE(op_to_count["Identity"] == 5); + ASSERT_TRUE(op_to_count["Dropout"] == 6); + + auto rule_transformer_L1 = std::make_unique("RuleTransformer1"); + rule_transformer_L1->Register(std::make_unique()); + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1); + ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1).IsOK()); + + op_to_count = CountOpsInGraph(graph); + // Of the 6 Dropout nodes in the graph, all but the ones named `d1` and `d6` should have been removed. + // A Dropout node can be removed if its second, optional output `mask` is either missing or unused downstream. + // `d1` cannot be removed because an Identity node has its `mask` output as an input; + // `d6` cannot be removed because its `mask` output is marked as a graph output. + ASSERT_TRUE(op_to_count["Identity"] == 5); + ASSERT_TRUE(op_to_count["Dropout"] == 2); +} + TEST(GraphTransformationTests, SliceElimination) { string model_uri = MODEL_FOLDER + "slice-elim.onnx"; std::shared_ptr model; diff --git a/onnxruntime/test/testdata/transform/dropout.onnx b/onnxruntime/test/testdata/transform/dropout.onnx new file mode 100644 index 0000000000000000000000000000000000000000..012cebe634a304ce723653ed6f45e57132573b10 GIT binary patch literal 718 zcmajbPfNov7{_rFD2orGGtP_0AvgpN{pZ2k+HT+>^eWyaoiNcw+U(%F@jdtk{k1)8 zO+#1c<;#aBc_KVH{?ukyJ(zN}t={!lSv~!EbVuw(vbkjHh8O;Hl$H8h@67!LdF9G% z^z{k8U&d(Gz(jK8Lk0u#wvj1RcvNLj(ST}CUEtIJGCnGokRjybfV?E^%X5D!o{RO? zY|IZKDWssAuxuW7Q;HOI6OPQ|eUhW1eZtqP`W?>cOfV^k3x=P+O$aADYtYnF%Z5%H yIcw~ys-18!_X8K3+q90ViIz-y5-X^71&gLVilLb+E_$&Z+IxZeQcIgAd`$oR0co88 literal 0 HcmV?d00001