ONNX: Pass name to the InputEdge (#12177)

This commit is contained in:
Artur Kulikowski 2022-07-19 08:52:54 +02:00 committed by GitHub
parent 6966af3ced
commit 6a795936b2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 126 additions and 1 deletions

View File

@ -0,0 +1,111 @@
ir_version: 8
producer_name: "test_data_generator"
graph {
node {
input: "input1"
output: "2872"
op_type: "Relu"
}
node {
input: "input2"
output: "2890"
op_type: "Relu"
}
node {
input: "2872"
input: "2890"
output: "2891"
name: "Add_221"
op_type: "Add"
}
node {
input: "2891"
output: "output0"
op_type: "Relu"
}
node {
input: "2891"
output: "output1"
op_type: "Relu"
}
name: "graph"
input {
name: "input1"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 4
}
dim {
dim_value: 512
}
}
}
}
}
input {
name: "input2"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 4
}
dim {
dim_value: 512
}
}
}
}
}
output {
name: "output0"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 4
}
dim {
dim_value: 512
}
}
}
}
}
output {
name: "output1"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 4
}
dim {
dim_value: 512
}
}
}
}
}
}
opset_import {
version: 16
}

View File

@ -206,7 +206,7 @@ std::vector<InputEdge> onnx_editor::EdgeMapper::find_output_consumers(const std:
const auto node_idx = it->second;
const auto port_indexes = get_node_input_indexes(node_idx, output_name);
for (const auto& idx : port_indexes) {
const auto consumer_edge = InputEdge{node_idx, idx};
const auto consumer_edge = InputEdge{node_idx, idx, output_name};
if (std::find_if(std::begin(input_edges), std::end(input_edges), [&consumer_edge](const InputEdge& edge) {
return edge.m_node_idx == consumer_edge.m_node_idx && edge.m_port_idx == consumer_edge.m_port_idx;
}) == std::end(input_edges)) {

View File

@ -1138,6 +1138,20 @@ NGRAPH_TEST(onnx_editor, editor_api_inputs_with_the_same_name) {
EXPECT_EQ(output_consumers[1].m_port_idx, 1);
}
NGRAPH_TEST(onnx_editor, editor_api_find_output_consumers_name) {
ONNXModelEditor editor{
ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests_3.onnx")};
const std::string output_name{"2891"};
std::vector<InputEdge> output_consumers = editor.find_output_consumers(output_name);
EXPECT_EQ(output_consumers[0].m_node_idx, 3);
EXPECT_EQ(output_consumers[0].m_port_idx, 0);
EXPECT_EQ(output_consumers[0].m_new_input_name, output_name);
EXPECT_EQ(output_consumers[1].m_node_idx, 4);
EXPECT_EQ(output_consumers[1].m_port_idx, 0);
EXPECT_EQ(output_consumers[1].m_new_input_name, output_name);
}
NGRAPH_TEST(onnx_editor, editor_api_is_correct_and_unambiguous_node) {
ONNXModelEditor editor{
ngraph::file_util::path_join(SERIALIZED_ZOO, "onnx/model_editor/subgraph_extraction_tests.onnx")};