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

Cleanup add/remove of initializer (#2274)

Browse source 
* Cleanup add/remove of initializer

* Address PR comments
This commit is contained in:
Scott McKay 2019-10-30 06:50:41 +10:00 committed by GitHub
parent 09eb8ff8b8
commit 47f40ca204
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 102 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
include/onnxruntime/core/graph/graph.h
View file

@ -966,7 +966,6 @@ class Graph {
ONNX_NAMESPACE::GraphProto* graph_proto_;
InitializedTensorSet name_to_initial_tensor_;
std::vector<int> removed_initializer_indexes_;
IOnnxRuntimeOpSchemaCollectionPtr schema_registry_;

15
onnxruntime/core/graph/function.cc
View file

@ -207,7 +207,11 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
for (const auto& input : meta_def->inputs) {
const ONNX_NAMESPACE::TensorProto* initializer = nullptr;
if (graph.GetInitializedTensor(input, initializer)) {
function_body_graph.AddInitializedTensor(*initializer);
// meta_def->inputs could have duplicates so make sure we only add once
const ONNX_NAMESPACE::TensorProto* subgraph_initializer = nullptr;
if (!function_body_graph.GetInitializedTensor(input, subgraph_initializer)) {
function_body_graph.AddInitializedTensor(*initializer);
}
}
}
@ -221,9 +225,9 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
const ONNX_NAMESPACE::FunctionProto& onnx_func_proto)
: parent_graph_(&graph) {
// Make a copy of the FunctionProto.
// All FunctionBody ops with the same op type seem to share the same FunctionProto struct within a model.
// All FunctionBody ops with the same op type seem to share the same FunctionProto struct within a model.
// Hence, we make a copy prior to generating the graph representation of the function,
// as we might make some modifications to the FunctionProto along the way
// as we might make some modifications to the FunctionProto along the way
onnx_func_proto_ = onnx_func_proto;
auto node_in_parent_graph = parent_graph_->GetNode(node_index);
@ -360,12 +364,13 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
new_attr_map[(*node_attr).name()] = *node_attr;
}
}
function_body_graph.AddNode(uniq_identifier + "_" + std::to_string(node_index), (*node).op_type(), (*node).doc_string(), inputs, outputs, &new_attr_map, (*node).domain());
function_body_graph.AddNode(uniq_identifier + "_" + std::to_string(node_index), (*node).op_type(),
(*node).doc_string(), inputs, outputs, &new_attr_map, (*node).domain());
}
auto status = function_body_graph.Resolve();
ORT_ENFORCE(status.IsOK(), "Resolve subgraph failed:", status.ErrorMessage());
} // namespace onnxruntime
}
FunctionImpl::~FunctionImpl() = default;

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

@ -2029,7 +2029,10 @@ void Graph::SetDescription(const std::string& description) {
}
void Graph::AddInitializedTensor(const TensorProto& tensor) {
if (name_to_initial_tensor_.end() != name_to_initial_tensor_.find(tensor.name())) {
auto existing = name_to_initial_tensor_.find(tensor.name());
if (existing != name_to_initial_tensor_.cend()) {
ORT_ENFORCE(existing->second == &tensor,
"AddInitializedTensor already has tensor with name ", tensor.name(), " but different TensorProto.");
return;
}
@ -2046,25 +2049,43 @@ void Graph::AddInitializedTensor(const TensorProto& tensor) {
ORT_IGNORE_RETURN_VALUE(GetOrCreateNodeArg(tensor.name(), &t));
}
SetGraphProtoSyncNeeded();
SetGraphResolveNeeded();
}
void Graph::RemoveInitializedTensor(const std::string& tensor_name) {
bool found = false;
auto iter = name_to_initial_tensor_.find(tensor_name);
if (name_to_initial_tensor_.end() != iter) {
found = iter != name_to_initial_tensor_.end();
if (found) {
name_to_initial_tensor_.erase(tensor_name);
SetGraphProtoSyncNeeded();
SetGraphResolveNeeded();
}
auto& mutable_initializers = *(graph_proto_->mutable_initializer());
auto proto_entry = std::find_if(mutable_initializers.begin(), mutable_initializers.end(),
[&tensor_name](const TensorProto& entry) { return entry.name() == tensor_name; });
if (proto_entry != mutable_initializers.end()) {
auto num_entries = mutable_initializers.size();
if (num_entries > 1) {
// swap the entry being deleted with the last one, and delete it.
// we do this so we don't have to move all the entries past the one being deleted down one.
auto slot = proto_entry - mutable_initializers.begin();
auto last_entry = mutable_initializers.end() - 1;
mutable_initializers.SwapElements(slot, num_entries - 1);
mutable_initializers.erase(last_entry);
} else {
mutable_initializers.erase(proto_entry);
}
} else {
// these should always be in sync as the pointer in name_to_initial_tensor_ is to memory owned by graph_proto_
ORT_ENFORCE(!found, "graph_proto_ is not in sync with name_to_initial_tensor_.");
}
}
Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_initializer) {
// name_to_initial_tensor_ maps from name to const TensorProto*, so we first
// look up the const pointer by name, then find and modify the mutable
// pointed-to TensorProto in graph_proto_.
const auto& initializer_name = new_initializer.name();
const auto name_to_initializer_it = name_to_initial_tensor_.find(initializer_name);
ORT_RETURN_IF_NOT(name_to_initializer_it != name_to_initial_tensor_.end(),
@ -2085,13 +2106,15 @@ Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_in
"Replacement tensor's data type does not match.");
auto& mutable_initializers = *(graph_proto_->mutable_initializer());
auto old_mutable_initializer_ptr_it = std::find(
mutable_initializers.pointer_begin(), mutable_initializers.pointer_end(), &old_initializer);
ORT_ENFORCE(old_mutable_initializer_ptr_it != mutable_initializers.pointer_end(),
"graph_proto_ is not in sync with name_to_initial_tensor_");
auto& old_mutable_initializer = **old_mutable_initializer_ptr_it;
// use cheaper pointer comparison to find old entry
auto existing_entry = std::find(mutable_initializers.pointer_begin(), mutable_initializers.pointer_end(),
&old_initializer);
old_mutable_initializer = new_initializer;
// these should always be in sync as the pointer in name_to_initial_tensor_ is to memory owned by graph_proto_
ORT_ENFORCE(existing_entry != mutable_initializers.pointer_end(),
"graph_proto_ is not in sync with name_to_initial_tensor_");
**existing_entry = new_initializer;
return Status::OK();
}
@ -2108,7 +2131,6 @@ bool Graph::GetInitializedTensor(const std::string& tensor_name, const TensorPro
void Graph::CleanAllInitializedTensors() noexcept {
name_to_initial_tensor_.clear();
removed_initializer_indexes_.clear();
// Clearing RepeatedPtrFields does not free objects' memory. The memory is retained
// and can be reused. Need to explicitly release the cleared objects and free the
@ -2288,35 +2310,6 @@ const ONNX_NAMESPACE::GraphProto& Graph::ToGraphProto() {
// Nodes.
ToGraphProtoInternal(*graph_proto_);
if (!removed_initializer_indexes_.empty()) {
// Move initializers.
std::sort(removed_initializer_indexes_.begin(), removed_initializer_indexes_.end());
int lastInUseInitializerIndex = graph_proto_->initializer_size() - 1;
int start = 0;
int end = gsl::narrow_cast<int>(removed_initializer_indexes_.size()) - 1;
int lastRemovedInitializerIndex = removed_initializer_indexes_[end];
for (; start <= end; start++) {
// Find a lastInUseInitializer.
while (start <= end && lastInUseInitializerIndex == lastRemovedInitializerIndex) {
graph_proto_->mutable_initializer()->RemoveLast();
lastInUseInitializerIndex--;
end--;
if (start <= end) {
lastRemovedInitializerIndex = removed_initializer_indexes_[end];
}
}
if (start <= end) {
// Copy the <lastInUseInitializerIndex> initializer in use to the <start> slot which is removed.
*graph_proto_->mutable_initializer(removed_initializer_indexes_[start]) = graph_proto_->initializer(lastInUseInitializerIndex);
graph_proto_->mutable_initializer()->RemoveLast();
lastInUseInitializerIndex--;
}
}
removed_initializer_indexes_.clear();
}
GraphProtoSyncNeeded(false);
return *graph_proto_;
@ -2329,9 +2322,7 @@ ONNX_NAMESPACE::GraphProto Graph::ToGraphProto() const {
GraphProto result;
ToGraphProtoInternal(result);
for (auto initializer : GetAllInitializedTensors()) {
*result.add_initializer() = *initializer.second;
}
*result.mutable_initializer() = graph_proto_->initializer();
return result;
}

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

@ -959,8 +959,7 @@ TEST(GraphUpdateTest, ReplaceInitializedTensor) {
status = graph.ReplaceInitializedTensor(valid_replacement);
ASSERT_TRUE(status.IsOK()) << status.ErrorMessage();
auto tensor_data_matches = [](
const ONNX_NAMESPACE::TensorProto& a, const ONNX_NAMESPACE::TensorProto& b) {
auto tensor_data_matches = [](const ONNX_NAMESPACE::TensorProto& a, const ONNX_NAMESPACE::TensorProto& b) {
if (a.int32_data_size() != b.int32_data_size()) return false;
for (int i = 0; i < a.int32_data_size(); ++i) {
if (a.int32_data(i) != b.int32_data(i)) return false;
@ -979,5 +978,59 @@ TEST(GraphUpdateTest, ReplaceInitializedTensor) {
ASSERT_TRUE(tensor_data_matches(graph_proto.initializer(0), valid_replacement));
}
}
TEST(GraphUpdateTest, AddRemoveInitializerHandling) {
Model m{"test_model"};
Graph& graph = m.MainGraph();
auto create_tensor_proto = [](const std::string& name, int32_t value) {
ONNX_NAMESPACE::TensorProto init{};
init.set_name(name);
init.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32);
init.add_dims(1);
init.add_int32_data(value);
return init;
};
auto init = create_tensor_proto("1", 1);
auto init2 = create_tensor_proto("2", 2);
// add both, remove the 1st (moves the second initializer into the first slot), and finally re-add the first
graph.AddInitializedTensor(init);
graph.AddInitializedTensor(init2);
graph.RemoveInitializedTensor(init.name());
graph.AddInitializedTensor(init);
ASSERT_EQ(graph.GetAllInitializedTensors().size(), 2);
// check the values coming from name_to_initial_tensor_ are good;
const TensorProto* i = nullptr;
ASSERT_TRUE(graph.GetInitializedTensor(init.name(), i));
ASSERT_TRUE(i->int32_data()[0] == 1);
ASSERT_TRUE(graph.GetInitializedTensor(init2.name(), i));
ASSERT_TRUE(i->int32_data()[0] == 2);
// check the values in the GraphProto are also correct
ONNX_NAMESPACE::GraphProto graph_proto_from_const_graph = static_cast<const Graph&>(graph).ToGraphProto();
ONNX_NAMESPACE::GraphProto graph_proto_from_graph = graph.ToGraphProto();
ASSERT_EQ(graph_proto_from_const_graph.initializer_size(), 2);
ASSERT_EQ(graph_proto_from_graph.initializer_size(), 2);
auto validate_proto = [&](const GraphProto& proto) {
auto initializers = proto.initializer();
// we expect '2' to be before '1' due to the remove moving the last initializer into the slot of the one being
// removed in order to free memory and only move one entry
EXPECT_EQ(initializers[0].name(), init2.name());
EXPECT_EQ(initializers[0].int32_data()[0], 2);
EXPECT_EQ(initializers[1].name(), init.name());
EXPECT_EQ(initializers[1].int32_data()[0], 1);
};
validate_proto(graph_proto_from_const_graph);
validate_proto(graph_proto_from_graph);
}
} // namespace test
} // namespace onnxruntime