saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support opset 11 subgraph of Squad model in Embed Layer Normalization (#2605)

Browse source 
Support opset 11 Squad model that is exported from PyTorch nightly. The embed layer uses Range op which is missed in the transformer.
This commit is contained in:
Tianlei Wu 2019-12-10 15:22:19 -08:00 committed by GitHub
parent 796948c6ae
commit bc89eccb21
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 307 additions and 77 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -65,6 +65,193 @@ static bool CheckInput(NodeArg* input, const logging::Logger& logger) {
return true;
}
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)) {
return false;
}
for (size_t i = 0; i < edges.size(); i++) {
if (edges[i]->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());
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) {
return false;
}
matched_edges.insert(matched_edges.end(), edges.begin(), edges.end());
return true;
}
/** Match subgraph like the following:
(input_ids)
/ \
Shape Shape
| |
Gather (indice=0) Gather (indice=1)--+
| | |
Unsqueeze Unsqueeze Cast
\ / |
\ / Range(start=0, delta=1)
\ / |
Concat Unsqueeze
| |
+--|----------------------------+
| |
Expand
|
Gather
Note that position gather node is the node in the bottom of above sub-graph.
*/
static bool MatchPositionEmbeddingSubgraph2(
Graph& graph,
Node& position_gather_node,
NodeArg* input_ids,
const logging::Logger& logger,
std::vector<const Node::EdgeEnd*>& matched_edges) {
// Match Gather <-- Expand <-- Unsqueeze <-- Range <-- Cast <-- Gather <-- Shape
std::vector<NodeIndex> position_parent_nodes;
std::vector<graph_utils::EdgeEndToMatch> position_embedding_path_symbolic{
{0, 1, "Expand", {8}, kOnnxDomain},
{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Range", {11}, kOnnxDomain},
{0, 1, "Cast", {9}, kOnnxDomain},
{0, 0, "Gather", {11}, kOnnxDomain},
{0, 0, "Shape", {1}, kOnnxDomain},
};
std::vector<const Node::EdgeEnd*> edges;
if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) {
DEBUG_LOG("Failed to find path 1.");
return false;
}
for (size_t i = 0; i < edges.size(); i++) {
if (edges[i]->GetNode().GetOutputEdgesCount() != (i == 4 ? 2 : 1)) {
DEBUG_LOG("Output edge count not expected for nodes in path 1.");
return false;
}
}
matched_edges = edges;
Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index());
Node& range_node = *graph.GetNode(edges[2]->GetNode().Index());
Node& gather_node_1 = *graph.GetNode(edges[4]->GetNode().Index());
Node& shape_node_1 = *graph.GetNode(edges[5]->GetNode().Index());
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(range_node.InputDefs()[0]), int64_t(0), true)) {
DEBUG_LOG("The first input of Range should be a constant with value 0.");
return false;
}
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(range_node.InputDefs()[2]), int64_t(1), true)) {
DEBUG_LOG("The third input of Range should be a constant with value 1.");
return false;
}
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node_1.InputDefs()[1]), int64_t(1), true)) {
DEBUG_LOG("The second input of Gather in path1 should be a constant with value 1.");
return false;
}
std::vector<graph_utils::EdgeEndToMatch> expand_parent_path1{
{0, 1, "Concat", {11}, kOnnxDomain},
{0, 0, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Gather", {1, 11}, kOnnxDomain},
{0, 0, "Shape", {1}, kOnnxDomain},
};
if (!graph_utils::FindPath(expand_node, true, expand_parent_path1, edges, logger)) {
DEBUG_LOG("Failed to find path 2.");
return false;
}
for (size_t i = 0; i < edges.size(); i++) {
if (edges[i]->GetNode().GetOutputEdgesCount() != 1) {
DEBUG_LOG("Output edge count not expected for nodes in path 2.");
return false;
}
}
Node& concat_node = *graph.GetNode(edges[0]->GetNode().Index());
Node& gather_node_0 = *graph.GetNode(edges[2]->GetNode().Index());
Node& shape_node_0 = *graph.GetNode(edges[3]->GetNode().Index());
if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node_0.InputDefs()[1]), int64_t(0), true)) {
DEBUG_LOG("Second input of Gather in path2 should be a constant with value 0.");
return false;
}
matched_edges.insert(matched_edges.end(), edges.begin(), edges.end());
std::vector<graph_utils::EdgeEndToMatch> concat_parent_path{
{0, 1, "Unsqueeze", {1, 11}, kOnnxDomain},
{0, 0, "Gather", {1, 11}, kOnnxDomain}
};
if (!graph_utils::FindPath(concat_node, true, concat_parent_path, edges, logger)) {
DEBUG_LOG("Failed to find path 3.");
return false;
}
// Two paths share the gather node (with second input indices==1)
if (edges[1]->GetNode().Index() != gather_node_1.Index()) {
DEBUG_LOG(" Gather nodes in path 1 and 3 expected to be same node.");
return false;
}
if (edges[0]->GetNode().GetOutputEdgesCount() != 1) {
DEBUG_LOG("Output edge count not expected for nodes in path 3.");
return false;
}
// Check if the two paths of position gather lead to the same input.
if (shape_node_0.MutableInputDefs()[0] != input_ids ||
shape_node_1.MutableInputDefs()[0] != input_ids) {
DEBUG_LOG("The parent of two shape nodes are expected to be input_ids.");
return false;
}
// Do not add the gather node since it has been added in another path.
matched_edges.push_back(edges[0]);
return true;
}
/**
Embed Layer Normalization will fuse embeddings and mask processing into one node :
The embeddings before conversion:
@ -176,22 +363,19 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l
continue;
}
NodeArg* input_ids = word_gather_node.MutableInputDefs()[1];
// Check the second input of position gather. If it's not initializer, check for two paths.
Node* p_expand_node = nullptr;
Node* p_shape_node = nullptr;
std::vector<const Node::EdgeEnd*> pg_edges;
bool isValidEmbedSubNode = true;
if (graph_utils::IsConstantInitializer(graph, position_gather_node.MutableInputDefs()[1]->Name())) {
// Check if the second input of position gather is a tensor with values evenly spaced by 1 starting from 0.
// Check if the second input of position gather is a tensor with values evenly spaced by 1 starting from 0.
std::vector<int64_t> data;
auto expected_shape = word_gather_node.MutableInputDefs()[1]->Shape();
if (!optimizer_utils::AppendTensorFromInitializer(graph, *(position_gather_node.MutableInputDefs()[1]), data)
|| !utils::HasDimValue(expected_shape->dim()[0])
|| !utils::HasDimValue(expected_shape->dim()[1])
|| static_cast<int>(data.size()) != expected_shape->dim()[0].dim_value() * expected_shape->dim()[1].dim_value()) {
if (!optimizer_utils::AppendTensorFromInitializer(graph, *(position_gather_node.MutableInputDefs()[1]), data) || !utils::HasDimValue(expected_shape->dim()[0]) || !utils::HasDimValue(expected_shape->dim()[1]) || static_cast<int>(data.size()) != expected_shape->dim()[0].dim_value() * expected_shape->dim()[1].dim_value()) {
continue;
}
int64_t expected_value = 0;
bool isValidEmbedSubNode = true;
for (size_t i = 0; i < data.size(); i++) {
if (data[i] != expected_value) {
isValidEmbedSubNode = false;
@ -202,68 +386,19 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l
expected_value = 0;
}
}
if (!isValidEmbedSubNode) {
continue;
}
} else {
// 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}};
if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) {
continue;
}
if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) {
continue;
}
p_expand_node = graph.GetNode(edges[0]->GetNode().Index());
p_shape_node = graph.GetNode(edges[1]->GetNode().Index());
// 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},
};
if (!graph_utils::FindPath(expand_node, true, pg_parent_path, pg_edges, logger)) {
continue;
}
for (size_t i = 0; i < pg_edges.size(); i++) {
if (pg_edges[i]->GetNode().GetOutputEdgesCount() != 1) {
isValidEmbedSubNode = false;
break;
if (!MatchPositionEmbeddingSubgraph1(graph, position_gather_node, input_ids, logger, pg_edges)) {
pg_edges.clear();
if (!MatchPositionEmbeddingSubgraph2(graph, position_gather_node, input_ids, logger, pg_edges)) {
continue;
}
}
// Check if the second input of the Gather node in the path has a constant input of 1
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. ");
continue;
}
// Check if the two paths of position gather lead to the same input.
Node& shape_node_2 = *graph.GetNode(pg_edges[pg_edges.size() - 1]->GetNode().Index());
if (shape_node_1.MutableInputDefs()[0] != shape_node_2.MutableInputDefs()[0]) {
continue;
}
// Check if the parent of "shape" is the parent of "word gather"
if (shape_node_1.MutableInputDefs()[0] != word_gather_node.MutableInputDefs()[1]) {
continue;
}
}
if (!isValidEmbedSubNode) {
continue;
}
// Get input "input_ids" from node.
NodeArg* input_ids = word_gather_node.MutableInputDefs()[1];
if (!CheckInput(input_ids, logger)) {
DEBUG_LOG("Input id is not valid. ");
continue;
@ -283,32 +418,30 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l
continue;
}
if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) !=
utils::GetTensorShapeFromTensorShapeProto(*(segment_ids->Shape()))) {
if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) !=
utils::GetTensorShapeFromTensorShapeProto(*(segment_ids->Shape()))) {
DEBUG_LOG("Input_ids and segment id should have the same shape. ");
continue;
}
if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) !=
utils::GetTensorShapeFromTensorShapeProto(*(mask->Shape()))) {
if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) !=
utils::GetTensorShapeFromTensorShapeProto(*(mask->Shape()))) {
DEBUG_LOG("Input_ids and mask should have the same shape. ");
continue;
}
NodeArg* gamma = layer_norm_node.MutableInputDefs()[1];
NodeArg* beta = layer_norm_node.MutableInputDefs()[2];
if (gamma->Shape() == nullptr
|| gamma->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) {
if (gamma->Shape() == nullptr || gamma->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) {
DEBUG_LOG("Gamma should be of shape (hidden_size). ");
continue;
}
if (beta->Shape() == nullptr
|| beta->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) {
if (beta->Shape() == nullptr || beta->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) {
DEBUG_LOG("Beta should be of shape (hidden_size). ");
continue;
}
// Cast input_ids, segment_ids, and mask to int32 if needed.
// Cast input_ids, segment_ids, and mask to int32 if needed.
input_ids = CastToInt32(graph, input_ids, layer_norm_node.GetExecutionProviderType());
segment_ids = CastToInt32(graph, segment_ids, layer_norm_node.GetExecutionProviderType());
mask = CastToInt32(graph, mask, layer_norm_node.GetExecutionProviderType());
@ -339,10 +472,6 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l
for (size_t i = 0; i < pg_edges.size(); i++) {
nodes_to_remove.push_back(pg_edges[i]->GetNode().Index());
}
if (p_shape_node != nullptr && p_expand_node != nullptr) {
nodes_to_remove.push_back(p_shape_node->Index());
nodes_to_remove.push_back(p_expand_node->Index());
}
nodes_to_remove.push_back(word_gather_node.Index());
nodes_to_remove.push_back(position_gather_node.Index());
nodes_to_remove.push_back(segment_gather_node.Index());

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

@ -1328,6 +1328,32 @@ TEST(GraphTransformationTests, EmbedLayerNormFusionFormat2) {
ASSERT_TRUE(op_to_count["SkipLayerNormalization"] == 0);
ASSERT_TRUE(op_to_count["EmbedLayerNormalization"] == 1);
}
TEST(GraphTransformationTests, EmbedLayerNormFusionFormat3) {
auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format3.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);
EXPECT_EQ(op_to_count["Shape"], 0);
EXPECT_EQ(op_to_count["Expand"], 0);
EXPECT_EQ(op_to_count["Gather"], 0);
EXPECT_EQ(op_to_count["Unsqueeze"], 0);
EXPECT_EQ(op_to_count["LayerNormalization"], 0);
EXPECT_EQ(op_to_count["SkipLayerNormalization"], 0);
EXPECT_EQ(op_to_count["ReduceSum"], 0);
EXPECT_EQ(op_to_count["MatMul"], 1);
EXPECT_EQ(op_to_count["Add"], 2);
EXPECT_EQ(op_to_count["Cast"], 3);
EXPECT_EQ(op_to_count["Attention"], 1);
EXPECT_EQ(op_to_count["EmbedLayerNormalization"], 1);
}
#endif
} // namespace test

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

Binary file not shown.

75
onnxruntime/test/testdata/transform/fusion/embed_layer_norm_gen.py vendored Normal file
View file

@ -0,0 +1,75 @@
import onnx
from onnx import helper
from onnx import TensorProto
from enum import Enum
class Format(Enum):
Format1=1,
Format2=2,
Format3=3
def GenerateModel(format, model_name):
nodes = [ # LayerNorm subgraph
helper.make_node("Shape", ["input_ids"], ["shape1_out"], "shape1"),
helper.make_node("Gather", ["shape1_out", "indices_0"], ["gather0_out"], "gather0"),
helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]),
helper.make_node("Shape", ["input_ids"], ["shape2_out"], "shape2"),
helper.make_node("Gather", ["shape2_out", "indices_1"], ["gather1_out"], "gather1"),
helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]),
helper.make_node("Concat", ["unsqueeze0_out", "unsqueeze1_out"], ["concat_out"], "concat", axis=0),
helper.make_node("Cast", ["gather1_out"], ["cast_out"], "cast", to=7),
helper.make_node("Range", ["start_0", "cast_out", "delta_1"], ["range_out"], "range"),
helper.make_node("Unsqueeze", ["range_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]),
helper.make_node("Expand", ["unsqueeze2_out", "concat_out"], ["expand_out"], "expand"),
helper.make_node("Gather", ["pos_embed", "expand_out"], ["pos_gather_out"], "pos_gather"),
helper.make_node("Gather", ["word_embed", "input_ids"], ["word_gather_out"], "word_gather"),
helper.make_node("Add", ["word_gather_out", "pos_gather_out"], ["word_add_pos_out"], "word_add_pos"),
helper.make_node("Gather", ["seg_embed", "segment_ids"], ["seg_gather_out"], "seg_gather"),
helper.make_node("Add", ["word_add_pos_out", "seg_gather_out"], ["add3_out"], "add3"),
helper.make_node("LayerNormalization", ["add3_out", "layer_norm_weight", "layer_norm_bias"], ["layernorm_out"], "layernorm", axis=-1, epsion=0.000009999999747378752),
helper.make_node("Cast", ["input_mask"], ["mask_cast_out"], "mask_cast", to=6),
helper.make_node("ReduceSum", ["mask_cast_out"], ["mask_index_out"], "mask_index", axes=[1], keepdims=0),
helper.make_node("Attention", ["layernorm_out", "qkv_weights", "qkv_bias", "mask_index_out"], ["att_out"], "att", domain="com.microsoft", num_heads=2),
helper.make_node("MatMul", ["att_out", "matmul_weight"], ["matmul_out"], "matmul"),
helper.make_node("Add", ["matmul_out", "add_bias"], ["add_out"], "add"),
helper.make_node("Add", ["add_out", "layernorm_out"], ["add2_out"], "add2")
]
# hidden_size=4, num_heads=2, max_seq_length=3
initializers = [ # initializers
helper.make_tensor('indices_0', TensorProto.INT64, [], [0]),
helper.make_tensor('indices_1', TensorProto.INT64, [], [1]),
helper.make_tensor('start_0', TensorProto.INT64, [], [0]),
helper.make_tensor('delta_1', TensorProto.INT64, [], [1]),
helper.make_tensor('word_embed', TensorProto.FLOAT, [2, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]),
helper.make_tensor('pos_embed', TensorProto.FLOAT, [4, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]),
helper.make_tensor('seg_embed', TensorProto.FLOAT, [2, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]),
helper.make_tensor('layer_norm_weight', TensorProto.FLOAT, [4], [1.0, 2.0, 3.0, 4.0]),
helper.make_tensor('layer_norm_bias', TensorProto.FLOAT, [4], [0.1, 0.2, 0.3, 0.4]),
helper.make_tensor('qkv_weights', TensorProto.FLOAT, [4, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]),
helper.make_tensor('qkv_bias', TensorProto.FLOAT, [4], [0.1, 0.2, 0.3, 0.4]),
helper.make_tensor('matmul_weight', TensorProto.FLOAT, [4, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]),
helper.make_tensor('add_bias', TensorProto.FLOAT, [4], [0.1, 0.2, 0.3, 0.4]),
]
graph = helper.make_graph(
nodes,
"EmbedLayerNorm_format3", #name
[ # inputs
helper.make_tensor_value_info('input_ids', TensorProto.INT64, ['batch', 3]),
helper.make_tensor_value_info('segment_ids', TensorProto.INT64, ['batch', 3]),
helper.make_tensor_value_info('input_mask', TensorProto.INT64, ['batch', 3]),
],
[ # outputs
helper.make_tensor_value_info('add2_out', TensorProto.FLOAT, ['batch', 3, 4]),
],
initializers
)
model = helper.make_model(graph)
onnx.save(model, model_name)
GenerateModel(Format.Format3, 'embed_layer_norm_format3.onnx')