saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Layer Norm Fusion Fix (#2379)

Browse source 
* layer norm fusion fix

* Add input shape check in code and unit tests
This commit is contained in:
liuziyue 2019-11-12 17:19:51 -08:00 committed by GitHub
parent 8c733c8d82
commit 4b72fedbd5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 44 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
onnxruntime/core/optimizer/layer_norm_fusion.cc
View file

@ -218,25 +218,36 @@ Status LayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level)
}
nodes_to_remove.push_back(last_add_node);
// Get the inputs for the new LayerNormalization node
// Get the inputs for the new LayerNormalization node.
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])) {
scale = mul_node.MutableInputDefs()[i];
// Scale must be 1d.
if (mul_node.MutableInputDefs()[i]->Shape()->dim_size() == 1) {
scale = mul_node.MutableInputDefs()[i];
}
}
}
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])) {
bias = mul_node.MutableInputDefs()[i];
// Bias must be 1d.
if (last_add_node.MutableInputDefs()[i]->Shape()->dim_size() == 1) {
bias = last_add_node.MutableInputDefs()[i];
}
}
}
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;
}
const std::vector<NodeArg*> layer_norm_input_defs{reduce_mean_node.MutableInputDefs()[0], scale, bias};
Node& layer_norm_node = graph.AddNode(graph.GenerateNodeName("LayerNormalization"),
"LayerNormalization",

30
onnxruntime/test/optimizer/graph_transform_test.cc
View file

@ -839,6 +839,21 @@ TEST(GraphTransformationTests, LayerNormFusionTest) {
ASSERT_TRUE(op_to_count["Pow"] == 0);
ASSERT_TRUE(op_to_count["Sqrt"] == 0);
ASSERT_TRUE(op_to_count["LayerNormalization"] == 1);
for (const Node& node : graph.Nodes()) {
if (node.OpType() == "LayerNormalization") {
// LayerNormalization should have three inputs.
EXPECT_EQ(node.InputDefs().size(), 3) << "LayerNormalization number of inputs does not equal to 3. Got:" << node.InputDefs().size();
// LayerNormalization input "scale" and "bias" should have the same dimension.
const TensorShapeProto* scale_shape = node.InputDefs()[1]->Shape();
const TensorShapeProto* bias_shape = node.InputDefs()[2]->Shape();
EXPECT_EQ(scale_shape->dim_size(), 1) << "LayerNormalization scale should be 1D. Got: " << scale_shape->dim_size();
EXPECT_EQ(bias_shape->dim_size(), 1) << "LayerNormalization bias should be 1D. Got: " << bias_shape->dim_size();
EXPECT_EQ(scale_shape->dim(0).dim_value(), bias_shape->dim(0).dim_value());
} else {
EXPECT_TRUE(false) << "Unexpected node " << node.Name();
}
}
}
TEST(GraphTransformationTests, LayerNormWithSubDupFusionTest) {
@ -860,6 +875,21 @@ TEST(GraphTransformationTests, LayerNormWithSubDupFusionTest) {
ASSERT_TRUE(op_to_count["Pow"] == 0);
ASSERT_TRUE(op_to_count["Sqrt"] == 0);
ASSERT_TRUE(op_to_count["LayerNormalization"] == 1);
for (const Node& node : graph.Nodes()) {
if (node.OpType() == "LayerNormalization") {
// LayerNormalization should have three inputs.
EXPECT_EQ(node.InputDefs().size(), 3) << "LayerNormalization number of inputs does not equal to 3. Got:" << node.InputDefs().size();
// LayerNormalization input "scale" and "bias" should have the same dimension.
const TensorShapeProto* scale_shape = node.InputDefs()[1]->Shape();
const TensorShapeProto* bias_shape = node.InputDefs()[2]->Shape();
EXPECT_EQ(scale_shape->dim_size(), 1) << "LayerNormalization scale should be 1D. Got: " << scale_shape->dim_size();
EXPECT_EQ(bias_shape->dim_size(), 1) << "LayerNormalization bias should be 1D. Got: " << bias_shape->dim_size();
EXPECT_EQ(scale_shape->dim(0).dim_value(), bias_shape->dim(0).dim_value());
} else {
EXPECT_TRUE(false) << "Unexpected node " << node.Name();
}
}
}
#endif