saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-06-04 23:59:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add check to prevent storing nullptr in value_info_ when proto has unused value info (#3461)

Browse source 
* Add unit test for serialization of unused value_info

* Do not add non-existent (nullptr) value_info_ when loading a model.

Fixes #3430
This commit is contained in:
Pranav Prakash 2020-04-09 19:25:10 -07:00 committed by GitHub
parent 2ccedb7b4d
commit 95ade8f47b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/core/graph/graph.cc
View file

@ -2530,7 +2530,9 @@ Status Graph::SetGraphInputsOutputs() {
for (auto& graph_value_info : graph_proto_->value_info()) {
auto& name = graph_value_info.name();
const auto* node_arg = GetNodeArg(name);
value_info_.push_back(node_arg);
if (node_arg != nullptr) {
value_info_.push_back(node_arg);
}
}
} else {

25
onnxruntime/test/ir/graph_test.cc
View file

@ -246,6 +246,31 @@ TEST_F(GraphTest, SimpleUnique) {
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(std::move(m), model, nullptr, *logger_));
}
TEST_F(GraphTest, UnusedValueInfoSerializes) {
ModelProto m;
m.set_ir_version(4);
ImportOpset(m, "", 11);
GraphProto& g = *m.mutable_graph();
NodeProto* node = g.add_node();
*node->add_input() = "x";
*node->add_output() = "sum";
node->set_op_type("Unique");
node->set_domain("");
ValueInfoProto* input1 = g.add_input();
input1->set_name("x");
SetTypeAndShape(input1->mutable_type()->mutable_tensor_type(), 1, {3, 4, 5});
ValueInfoProto* output = g.add_output();
output->set_name("sum");
SetTypeAndShape(output->mutable_type()->mutable_tensor_type(), 1, {60});
ValueInfoProto* unused = g.add_value_info();
unused->set_name("unused");
SetTypeAndShape(unused->mutable_type()->mutable_tensor_type(), 1, {123});
std::shared_ptr<Model> model;
ASSERT_STATUS_OK(Model::Load(std::move(m), model, nullptr, *logger_));
model->MainGraph().SetGraphProtoSyncNeeded();
EXPECT_TRUE(Model::Save(*model, "graph_with_unused_value_info.onnx").IsOK());
}
TEST_F(GraphTest, WrongOpset) {
ModelProto m;