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

EmbedLayerNormalization Fusion For Dynamic Squad Model Opset 10 (#2613)

Browse source 
Support subgraph of SQuAD model exported from pytorch with dynamic input axes
This commit is contained in:
liuziyue 2019-12-10 23:40:55 -08:00 committed by Tianlei Wu
parent 2ca9733cee
commit 1ee250def8
3 changed files with 148 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

157
onnxruntime/core/optimizer/embed_layer_norm_fusion.cc
View file

@ -65,65 +65,146 @@ static bool CheckInput(NodeArg* input, const logging::Logger& logger) {
return true;
}
/** Match subgraph like the following:
(input_ids)
/ \
Shape Shape
| |
^Gather (indice=0)^ Gather (indice=1)--+
^|^ ^|^ |
^Unsqueeze^ ^Unsqueeze^ Unsqueeze
^\^ ^/^ |
^\^ ^/^ ConstantOfShape
^\^ ^/^ |
^Concat^ NonZero
| |
| Transpose
| |
| Squeeze
| |
| Cast
| |
| Unsqueeze
+--|----------------------------+
| |
Expand
|
Gather
Note that position gather node is the node in the bottom of above sub-graph.
Paths in ^^ are alternative path to be matched if path input_ids -> Shape -> Expand -> Gather is not found.
*/
static bool MatchPositionEmbeddingSubgraph1(
Graph& graph,
Node& position_gather_node,
NodeArg* input_ids,
const logging::Logger& logger,
std::vector<const Node::EdgeEnd*>& matched_edges) {
// Match two paths.
// Match Shape --> Expand path if needed.
std::vector<NodeIndex> position_parent_nodes;
std::vector<graph_utils::EdgeEndToMatch> position_embedding_path_symbolic{
{0, 1, "Expand", {8}, kOnnxDomain},
{0, 1, "Shape", {1}, kOnnxDomain}};
std::vector<const Node::EdgeEnd*> edges;
if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) {
return false;
}
if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) {
return false;
}
// Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand
Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index());
Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index());
std::vector<graph_utils::EdgeEndToMatch> pg_parent_path{
{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Cast", {9}, kOnnxDomain},
{0, 0, "Squeeze", {1}, kOnnxDomain},
{0, 0, "Transpose", {1}, kOnnxDomain},
{0, 0, "NonZero", {9}, kOnnxDomain},
{0, 0, "ConstantOfShape", {9}, kOnnxDomain},
{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Gather", {1, 11}, kOnnxDomain},
{0, 0, "Shape", {1}, kOnnxDomain},
};
matched_edges = edges;
if (!graph_utils::FindPath(expand_node, true, pg_parent_path, edges, logger)) {
std::vector<const Node::EdgeEnd*> pg_edges;
// Find the "Expand" node
if (!graph_utils::FindPath(position_gather_node, true, {{0, 1, "Expand", {8}, kOnnxDomain}}, pg_edges, logger)) {
return false;
}
for (size_t i = 0; i < edges.size(); i++) {
if (edges[i]->GetNode().GetOutputEdgesCount() != 1) {
Node& expand_node = *graph.GetNode(pg_edges[0]->GetNode().Index());
const Node::EdgeEnd* expand_edge = pg_edges[0];
// Look for Path 1:
// Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand
if (!graph_utils::FindPath(expand_node, true,
{{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Cast", {9}, kOnnxDomain},
{0, 0, "Squeeze", {1}, kOnnxDomain},
{0, 0, "Transpose", {1}, kOnnxDomain},
{0, 0, "NonZero", {9}, kOnnxDomain},
{0, 0, "ConstantOfShape", {9}, kOnnxDomain},
{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Gather", {1, 11}, kOnnxDomain},
{0, 0, "Shape", {1}, kOnnxDomain}},
pg_edges, logger)) {
return false;
}
// All nodes in Path 1 except the "Gather" node must have only 1 output edge.
for (size_t i = 0; i < pg_edges.size() - 2; i++) {
if (pg_edges[i]->GetNode().GetOutputEdgesCount() != 1) {
return false;
}
}
if (pg_edges[8]->GetNode().GetOutputEdgesCount() != 1) {
return false;
}
// Check if the second input of the Gather node in the path has a constant input of 1
Node& gather_node = *graph.GetNode(edges[edges.size() - 2]->GetNode().Index());
Node& gather_node = *graph.GetNode(pg_edges[pg_edges.size() - 2]->GetNode().Index());
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) {
DEBUG_LOG("Second input of Gather should be a constant with value 1. ");
return false;
}
// Check if the parent of "shape" is the input_ids
Node& shape_node_2 = *graph.GetNode(edges[edges.size() - 1]->GetNode().Index());
if (shape_node_1.MutableInputDefs()[0] != input_ids ||
shape_node_2.MutableInputDefs()[0] != input_ids) {
// Match Shape --> Expand path if needed.
std::vector<const Node::EdgeEnd*> pg_edges_2;
Node* p_shape_node_2 = nullptr;
const Node::EdgeEnd* unsqueeze_edge = nullptr;
if (pg_edges[7]->GetNode().GetOutputEdgesCount() == 1) {
// In this case, the "Gather" node in "Path 1" must have 1 output edge.
if (!graph_utils::FindPath(expand_node, true, {{0, 1, "Shape", {1}, kOnnxDomain}}, pg_edges_2, logger)) {
return false;
}
p_shape_node_2 = graph.GetNode(pg_edges_2[0]->GetNode().Index());
} else if (pg_edges[7]->GetNode().GetOutputEdgesCount() == 2) {
// In this case, the "Gather" node in "Path 1" must have 2 output edges.
if (!graph_utils::FindPath(expand_node, true,
{{0, 1, "Concat", {4, 11}, kOnnxDomain},
{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Gather", {1, 11}, kOnnxDomain},
{0, 0, "Shape", {1}, kOnnxDomain}},
pg_edges_2, logger)) {
return false;
}
p_shape_node_2 = graph.GetNode(pg_edges_2[3]->GetNode().Index());
// Check for Unsqueeze --> Concat
Node& concat_node = *graph.GetNode(pg_edges_2[0]->GetNode().Index());
std::vector<const Node::EdgeEnd*> pg_edges_3;
if (!graph_utils::FindPath(concat_node, true,
{{0, 1, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Gather", {1, 11}, kOnnxDomain}},
pg_edges_3, logger)) {
return false;
}
unsqueeze_edge = pg_edges_3[0];
if (pg_edges_3[0]->GetNode().GetOutputEdgesCount() != 1 || pg_edges_3[1]->GetNode().GetOutputEdgesCount() != 2) {
return false;
}
// The gather node must be the same gather node in path 1.
if (graph.GetNode(pg_edges_3[1]->GetNode().Index()) != graph.GetNode(pg_edges[7]->GetNode().Index())) {
return false;
}
} else {
return false;
}
for (size_t i = 0; i < pg_edges_2.size(); i++) {
if (pg_edges_2[i]->GetNode().GetOutputEdgesCount() != 1) {
return false;
}
}
// Check if the two paths of position gather lead to the same input.
Node& shape_node_1 = *graph.GetNode(pg_edges[pg_edges.size() - 1]->GetNode().Index());
Node& shape_node_2 = *graph.GetNode(p_shape_node_2->Index());
if (shape_node_1.MutableInputDefs()[0] != shape_node_2.MutableInputDefs()[0]) {
return false;
}
// Check if the parent of "shape" is the parent of "word gather"
if (shape_node_1.MutableInputDefs()[0] != input_ids) {
return false;
}
matched_edges.insert(matched_edges.end(), edges.begin(), edges.end());
// Add all the valid nodes to result.
matched_edges = pg_edges;
matched_edges.push_back(expand_edge);
matched_edges.insert(matched_edges.end(), pg_edges_2.begin(), pg_edges_2.end());
if (unsqueeze_edge != nullptr) {
matched_edges.push_back(unsqueeze_edge);
}
return true;
}

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

@ -1354,6 +1354,35 @@ TEST(GraphTransformationTests, EmbedLayerNormFusionFormat3) {
EXPECT_EQ(op_to_count["Attention"], 1);
EXPECT_EQ(op_to_count["EmbedLayerNormalization"], 1);
}
TEST(GraphTransformationTests, EmbedLayerNormFusionFormat4) {
auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format4.onnx";
std::shared_ptr<Model> p_model;
ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, DefaultLoggingManager().DefaultLogger()).IsOK());
Graph& graph = p_model->MainGraph();
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(onnxruntime::make_unique<EmbedLayerNormFusion>(), TransformerLevel::Level2);
auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger());
ASSERT_TRUE(ret.IsOK());
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Shape"] == 0);
ASSERT_TRUE(op_to_count["Expand"] == 0);
ASSERT_TRUE(op_to_count["Gather"] == 0);
ASSERT_TRUE(op_to_count["Concat"] == 0);
ASSERT_TRUE(op_to_count["Unsqueeze"] == 0);
ASSERT_TRUE(op_to_count["ConstantOfShape"] == 0);
ASSERT_TRUE(op_to_count["NonZero"] == 0);
ASSERT_TRUE(op_to_count["Transpose"] == 0);
ASSERT_TRUE(op_to_count["Squeeze"] == 0);
ASSERT_TRUE(op_to_count["Add"] == 0);
ASSERT_TRUE(op_to_count["ReduceSum"] == 0);
ASSERT_TRUE(op_to_count["Attention"] == 1);
ASSERT_TRUE(op_to_count["SkipLayerNormalization"] == 0);
ASSERT_TRUE(op_to_count["EmbedLayerNormalization"] == 1);
}
#endif
} // namespace test

BIN
onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format4.onnx vendored Normal file
View file

Binary file not shown.