From 381f4c442a6131bed5a8eff1951f3b2a865d93cd Mon Sep 17 00:00:00 2001 From: ytaous <4484531+ytaous@users.noreply.github.com> Date: Fri, 26 Jun 2020 00:04:12 -0700 Subject: [PATCH] LayerNormFusion - Cast support (#4320) * cast support for layernormfusion * cast support for layernormfusion * bug fix * fix build * bug fix * fix test * minor refactor * on comments * on comments * on comments * on comments Co-authored-by: Ethan Tao --- .../core/optimizer/layer_norm_fusion.cc | 74 +++++++++++++++--- .../test/optimizer/graph_transform_test.cc | 22 ++++++ .../fusion/layer_norm_with_cast.onnx | Bin 0 -> 621 bytes 3 files changed, 85 insertions(+), 11 deletions(-) create mode 100644 onnxruntime/test/testdata/transform/fusion/layer_norm_with_cast.onnx diff --git a/onnxruntime/core/optimizer/layer_norm_fusion.cc b/onnxruntime/core/optimizer/layer_norm_fusion.cc index f2ce33b645..fd82270923 100644 --- a/onnxruntime/core/optimizer/layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/layer_norm_fusion.cc @@ -45,6 +45,16 @@ X --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul | ^ | | +---------------------+ + +In recent pytorch, Cast nodes may be inserted before Pow to ensure that both inputs 'base' and 'power' are the same type +due to restriction in older opsets. Therefore, Layer Normalization will also handle the case below : ++---------------------+ +| | +| v +X --> ReduceMean --> Sub --> Cast --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add + | ^ + | | + +--------------------------------------------------------+ */ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const { GraphViewer graph_viewer(graph); @@ -76,7 +86,6 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, if ((*iter).OpType().compare("Sub") == 0) { if (subCnt == 0) { p_sub_node = &(*iter); - } else { p_sub_node_dup = &(*iter); } @@ -95,7 +104,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, Node& sub_node = *graph.GetNode(p_sub_node->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(sub_node, "Sub", {7}) || sub_node.GetExecutionProviderType() != reduce_mean_node.GetExecutionProviderType() || - !optimizer_utils::CheckOutputEdges(graph, sub_node, subCnt == 1 ? 2u: 1u) || + !optimizer_utils::CheckOutputEdges(graph, sub_node, subCnt == 1 ? 2u : 1u) || !IsSupportedDataType(sub_node)) { continue; } @@ -174,7 +183,7 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, continue; } nodes_to_remove.push_back(reduce_mean2_node); - + // Traceback the reduceMean node to find pow --> reduceMean Node& pow_node = *graph.GetNode(reduce_mean2_node.InputNodesBegin()->Index()); if (!graph_utils::IsSupportedOptypeVersionAndDomain(pow_node, "Pow", {7, 12}) || @@ -185,8 +194,21 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, } nodes_to_remove.push_back(pow_node); - // Traceback the pow node to find sub --> pow - const Node* p_sub2_node = graph_utils::FirstParentByType(pow_node, "Sub"); + // check if Cast node exists between sub and pow + const Node* p_cast_node = graph_utils::FirstParentByType(pow_node, "Cast"); + if (p_cast_node != nullptr) { + Node& cast_node = *graph.GetNode(pow_node.InputNodesBegin()->Index()); + if (!graph_utils::IsSupportedOptypeVersionAndDomain(cast_node, "Cast", {9}) || + cast_node.GetExecutionProviderType() != cast_node.GetExecutionProviderType() || + !optimizer_utils::CheckOutputEdges(graph, cast_node, 1) || + !IsSupportedDataType(cast_node)) { + continue; + } + nodes_to_remove.push_back(cast_node); + } + + // Traceback from the last node in vector to find sub --> pow or sub --> cast + const Node* p_sub2_node = graph_utils::FirstParentByType(nodes_to_remove.back(), "Sub"); if (p_sub2_node == nullptr || (p_sub2_node != p_sub_node && p_sub2_node != p_sub_node_dup)) { continue; @@ -212,36 +234,66 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, } nodes_to_remove.push_back(last_add_node); + // get axes attributes + const onnxruntime::NodeAttributes& attributes = reduce_mean_node.GetAttributes(); + std::vector axes_values; + if (attributes.find("axes") != attributes.end()) { + axes_values = RetrieveValues(attributes.at("axes")); + } + // Get the inputs for the new LayerNormalization node. + // scale and bias could be multi-dims; we only support it for training at the moment + // because SkipLayerNorm kernel, for example, has dependency on single dim size NodeArg* scale = nullptr; NodeArg* bias = nullptr; for (size_t i = 0; i < mul_node.MutableInputDefs().size(); i++) { if (graph_utils::NodeArgIsConstant(graph, *(mul_node.MutableInputDefs()[i])) || graph_utils::IsGraphInput(graph, mul_node.MutableInputDefs()[i])) { +#ifdef ENABLE_TRAINING + if (axes_values.empty() || + mul_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast(axes_values.size())) { + scale = mul_node.MutableInputDefs()[i]; + } +#else // Scale must be 1d. if (mul_node.MutableInputDefs()[i]->Shape()->dim_size() == 1) { scale = mul_node.MutableInputDefs()[i]; } +#endif } } for (size_t i = 0; i < last_add_node.MutableInputDefs().size(); i++) { if (graph_utils::NodeArgIsConstant(graph, *(last_add_node.MutableInputDefs()[i])) || graph_utils::IsGraphInput(graph, last_add_node.MutableInputDefs()[i])) { +#ifdef ENABLE_TRAINING + if (axes_values.empty() || + last_add_node.MutableInputDefs()[i]->Shape()->dim_size() == static_cast(axes_values.size())) { + bias = last_add_node.MutableInputDefs()[i]; + } +#else // Bias must be 1d. if (last_add_node.MutableInputDefs()[i]->Shape()->dim_size() == 1) { bias = last_add_node.MutableInputDefs()[i]; } +#endif } } if (scale == nullptr || bias == nullptr) { continue; } - // Scale and bias must have the same dimension. - if (scale->Shape()->dim(0).dim_value() != bias->Shape()->dim(0).dim_value()) { - continue; + // Scale and bias must have the same shape. + bool same_dim = true; + for (int i = 0; i < scale->Shape()->dim_size(); i++) { + if (scale->Shape()->dim(i).dim_value() != bias->Shape()->dim(i).dim_value()) { + same_dim = false; + break; + } } + if (!same_dim) + continue; + const std::vector layer_norm_input_defs{reduce_mean_node.MutableInputDefs()[0], scale, bias}; Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName("LayerNormalization"), "LayerNormalization", @@ -252,9 +304,9 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, // Get constant "epsilon" from "Add2" node if available. Else, default value will be used. const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_utils::GetConstantInitializer(graph, add2_node.MutableInputDefs()[1]->Name()); if (tensor_proto != nullptr && - tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { - Initializer initializer{*tensor_proto, graph.ModelPath()}; - layer_norm_node.AddAttribute("epsilon", initializer.data()[0]); + tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + Initializer initializer{*tensor_proto, graph.ModelPath()}; + layer_norm_node.AddAttribute("epsilon", initializer.data()[0]); } else { layer_norm_node.AddAttribute("epsilon", DEFAULT_LAYERNORM_EPSILON); } diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 517d381fdc..94be6c3469 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -2123,6 +2123,28 @@ TEST_F(GraphTransformationTests, LayerNormFusionTest) { } } +TEST_F(GraphTransformationTests, LayerNormWithCastFusionTest) { + auto model_uri = MODEL_FOLDER "fusion/layer_norm_with_cast.onnx"; + std::shared_ptr p_model; + ASSERT_STATUS_OK(Model::Load(model_uri, p_model, nullptr, *logger_)); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, *logger_); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + +#ifdef ENABLE_TRAINING + ASSERT_TRUE(op_to_count["Cast"] == 0); + ASSERT_TRUE(op_to_count["LayerNormalization"] == 1); +#else + ASSERT_TRUE(op_to_count["Cast"] == 1); + ASSERT_TRUE(op_to_count["LayerNormalization"] == 0); +#endif +} + TEST_F(GraphTransformationTests, LayerNormWithSubDupFusionTest) { auto model_uri = MODEL_FOLDER "fusion/layer_norm_sub_dup.onnx"; std::shared_ptr p_model; diff --git a/onnxruntime/test/testdata/transform/fusion/layer_norm_with_cast.onnx b/onnxruntime/test/testdata/transform/fusion/layer_norm_with_cast.onnx new file mode 100644 index 0000000000000000000000000000000000000000..85e71c2b749d4c42c42684e042867cf7e72ce6dd GIT binary patch literal 621 zcmb7C%SyvQ6z$Cbz^;P3WDg)&5+ulkE`cK^(Z8JLhm_&O98OFO8mbHom?U%-@an-%UJ4D;}lCud1K zRCS5pGj9XUc$!Q)AD{C9okjNvFR+Ab^JNeM4xz9v^Pci0Bft%WkV7Qw<9MP~Lz)G8 zWi)hFC9i=i2oV~7#Gxgu^XySIB_qH?gcgVPk66L~V%53f*99JD!-OwyZMO==*lvwWc1k`Z8y)Bw$18c#?Ig_X(M}}Me0xzML^#A|> literal 0 HcmV?d00001