Dropout op elimination - enable for ORT training (#5588)

* dropout elimination

* per comments

* fix build

* fix build

Co-authored-by: Ethan Tao <ettao@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>

onnxruntime/core/optimizer/dropout_elimination.cc

@@ -7,6 +7,7 @@
 #include "core/graph/graph_utils.h"
 #include "core/optimizer/rewrite_rule.h"
 #include "core/optimizer/dropout_elimination.h"
+#include "core/optimizer/initializer.h"
 
 namespace onnxruntime {
 
@@ -25,6 +26,44 @@ bool EliminateDropout::SatisfyCondition(const Graph& graph, const Node& node, co
     return false;
   }
 
+#ifdef ENABLE_TRAINING
+  // allow Dropout elimination when:
+  //    1. ratio input is an initializer of 0
+	//    2. ratio input is not a graph input, so it cannot be overridden
+
+  // support opset 12 and above for ort training
+  if (graph_utils::MatchesOpSinceVersion(node, {12}) && node.InputDefs().size() > 1) {
+    if (graph_utils::IsGraphInput(graph, node.InputDefs()[1])) {
+      return false;
+    }
+    const ONNX_NAMESPACE::TensorProto* initializer = graph.GetConstantInitializer(node.InputDefs()[1]->Name(), true);
+    if (!initializer) {
+      return false;
+    }
+    int32_t data_type = initializer->data_type();
+    Initializer ratio(*initializer, graph.ModelPath());
+    switch (data_type) {
+      case ONNX_NAMESPACE::TensorProto_DataType_FLOAT:
+        if (*ratio.data<float>() > 0.f) {
+          return false;
+        }
+        break;
+      case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
+        if (math::halfToFloat(ratio.data<MLFloat16>()->val) > 0.f) {
+          return false;
+        }
+        break;
+      case ONNX_NAMESPACE::TensorProto_DataType_DOUBLE:
+        if (*ratio.data<double>() > 0.f) {
+          return false;
+        }
+        break;
+      default:
+        ORT_THROW("Unexpected data type for Dropout 'ratio' input of ", initializer->data_type());
+    }
+  }
+#endif
+
   // A Dropout Node has one required output and an optional second output 'mask' at index == 1.
   // It can be safely removed if it has only one output that is used (checked by CanRemoveNode)
   // and that output is not the 'mask' output.

onnxruntime/test/testdata/transform/dropout_zeroratio_elimination.py

@@ -0,0 +1,60 @@
+import onnx
+from onnx import helper
+from onnx import TensorProto, OperatorSetIdProto
+
+# inputs/outputs
+X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [2, 1])
+O1 = helper.make_tensor_value_info('O1', TensorProto.FLOAT, [2, 1])
+O2 = helper.make_tensor_value_info('O2', TensorProto.FLOAT, [2, 1])
+O3 = helper.make_tensor_value_info('O3', TensorProto.FLOAT, [2, 1])
+O4 = helper.make_tensor_value_info('O4', TensorProto.FLOAT, [2, 1])
+O5 = helper.make_tensor_value_info('O5', TensorProto.FLOAT, [2, 1])
+
+X2 = helper.make_tensor_value_info('X2', TensorProto.FLOAT, [])
+
+# initializers
+zeroratio_float = helper.make_tensor('ratio_zero_float', TensorProto.FLOAT, [], [0.0])
+zeroratio_double = helper.make_tensor('ratio_zero_double', TensorProto.DOUBLE, [], [0.0])
+zeroratio_float16 = helper.make_tensor('ratio_zero_float16', TensorProto.FLOAT16, [], [0])
+nonzeroratio = helper.make_tensor('ratio_nonzero', TensorProto.FLOAT, [], [0.1])
+training_mode = helper.make_tensor('training_mode', TensorProto.BOOL, [], [1])
+
+opsets = []
+onnxdomain = OperatorSetIdProto()
+onnxdomain.version = 12
+onnxdomain.domain = "" # The empty string ("") or absence of this field implies the operator set that is defined as part of the ONNX specification.
+opsets.append(onnxdomain)
+
+kwargs={}
+kwargs['opset_imports'] = opsets
+
+# Create the model (ModelProto)
+I1 = helper.make_node('Identity', ['X'], ['I1_out'], name='I1')
+D1 = helper.make_node("Dropout", ["I1_out", "ratio_zero_float", "training_mode"], ["D1_out"], "D1")
+I2 = helper.make_node('Identity', ['D1_out'], ['O1'], name='I2')
+
+I3 = helper.make_node('Identity', ['X'], ['I3_out'], name='I3')
+D2 = helper.make_node("Dropout", ["I3_out", "ratio_nonzero", "training_mode"], ["D2_out"], "D2")
+I4 = helper.make_node('Identity', ['D2_out'], ['O2'], name='I4')
+
+I5 = helper.make_node('Identity', ['X'], ['I5_out'], name='I5')
+D3 = helper.make_node("Dropout", ["I5_out", "X2", "training_mode"], ["D3_out"], "D3")
+I6 = helper.make_node('Identity', ['D3_out'], ['O3'], name='I6')
+
+I7 = helper.make_node('Identity', ['X'], ['I7_out'], name='I7')
+D4 = helper.make_node("Dropout", ["I7_out", "ratio_zero_double", "training_mode"], ["D4_out"], "D4")
+I8 = helper.make_node('Identity', ['D4_out'], ['O4'], name='I8')
+
+I9 = helper.make_node('Identity', ['X'], ['I9_out'], name='I9')
+D5 = helper.make_node("Dropout", ["I9_out", "ratio_zero_float16", "training_mode"], ["D5_out"], "D5")
+I10 = helper.make_node('Identity', ['D5_out'], ['O5'], name='I10')
+
+graph = helper.make_graph(
+    [I1, D1, I2, I3, D2, I4, I5, D3, I6, I7, D4, I8, I9, D5, I10],
+    "Dropout_Elimination",  #name
+    [X, X2],
+    [O1, O2, O3, O4, O5],
+    [zeroratio_float, zeroratio_double, zeroratio_float16, nonzeroratio, training_mode])
+
+model = helper.make_model(graph, producer_name='onnx-example', **kwargs)
+onnx.save(model, 'dropout_ratio.onnx')

orttraining/orttraining/core/optimizer/graph_transformer_utils.cc

@@ -73,6 +73,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
       rule_transformer->Register(make_unique<UnsqueezeElimination>());
       rule_transformer->Register(make_unique<ExpandElimination>());
       rule_transformer->Register(make_unique<CastElimination>());
+      rule_transformer->Register(make_unique<EliminateDropout>());
       rule_transformer->Register(make_unique<NonZeroShapeSetter>());
       rule_transformer->Register(make_unique<InsertSoftmaxCrossEntropyLossOutput>());
 

orttraining/orttraining/test/optimizer/graph_transform_test.cc

@@ -10,6 +10,7 @@
 #include "gtest/gtest.h"
 #include "core/optimizer/rule_based_graph_transformer.h"
 #include "core/optimizer/utils.h"
+#include "core/optimizer/dropout_elimination.h"
 #include "orttraining/core/optimizer/bias_dropout_fusion.h"
 #include "orttraining/core/optimizer/gist_encode_decode.h"
 #include "orttraining/core/optimizer/nonzero_shape_setter.h"
@@ -30,6 +31,27 @@ namespace test {
 
 #define MODEL_FOLDER ORT_TSTR("testdata/transform/")
 
+TEST_F(GraphTransformationTests, DropoutWithZeroRatioElimination) {
+  auto model_uri = MODEL_FOLDER "dropout_ratio.onnx";
+  std::shared_ptr<Model> model;
+  ASSERT_STATUS_OK(Model::Load(model_uri, model, nullptr, *logger_));
+  Graph& graph = model->MainGraph();
+  std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
+  ASSERT_TRUE(op_to_count["Identity"] == 10);
+  ASSERT_TRUE(op_to_count["Dropout"] == 5);
+
+  auto rule_transformer_L1 = onnxruntime::make_unique<RuleBasedGraphTransformer>("RuleTransformer1");
+  rule_transformer_L1->Register(onnxruntime::make_unique<EliminateDropout>());
+  onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
+  graph_transformation_mgr.Register(std::move(rule_transformer_L1), TransformerLevel::Level1);
+  ASSERT_STATUS_OK(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1, *logger_));
+
+  op_to_count = CountOpsInGraph(graph);
+
+  ASSERT_TRUE(op_to_count["Identity"] == 10);
+  ASSERT_TRUE(op_to_count["Dropout"] == 2);
+}
+
 TEST_F(GraphTransformationTests, GistEncodeDecode) {
   auto model_uri = MODEL_FOLDER "../test_training_model.onnx";
   std::shared_ptr<Model> p_model;