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

@@ -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;
   }
 

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<Status(Graph&)> f
 Status ForAllSubgraphs(const Graph& main_graph, std::function<Status(const Graph&)> 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

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

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<std::string> 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

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<Node*> 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<const NodeArg*> 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

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> model;
+  ASSERT_TRUE(Model::Load(model_uri, model).IsOK());
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> 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<RuleBasedGraphTransformer>("RuleTransformer1");
+  rule_transformer_L1->Register(std::make_unique<EliminateDropout>());
+  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> model;