From 60bbdf14035014d94d525c402d3d6404ac32ad3d Mon Sep 17 00:00:00 2001 From: Guoyu Wang <62914304+gwang-msft@users.noreply.github.com> Date: Fri, 1 Oct 2021 11:44:26 -0700 Subject: [PATCH] Remove unused NodeArgs in Graph::Resolve (#9213) * Remove unused NodeArgs * Handle case where a node arg from an initializer from initializer_names_to_preserve * Fix CI failure * update test * Fix outer scope node args failure * Use NodeArg* as the key of the std::set instead of string * Minor updates --- include/onnxruntime/core/graph/graph.h | 12 ++--- onnxruntime/core/graph/graph.cc | 68 +++++++++++++++++++++----- onnxruntime/test/ir/graph_test.cc | 28 ++++++++--- 3 files changed, 82 insertions(+), 26 deletions(-) diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 898e2e5af0..490e4f13dc 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -644,7 +644,7 @@ class Graph { bool IsInitializedTensor(const std::string& name) const; #if !defined(DISABLE_SPARSE_TENSORS) - /** Check if a given name is a sparse initializer's name in the model + /** Check if a given name is a sparse initializer's name in the model * we currently convert sparse_initializer field in the model into dense Tensor instances. * However, we sometimes want to check if this initializer was stored as sparse in the model. */ @@ -674,7 +674,7 @@ class Graph { */ const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const std::string& name, bool check_outer_scope) const; - /** returns the initializer's TensorProto if 'name' is an initializer (both constant and overridable). + /** returns the initializer's TensorProto if 'name' is an initializer (both constant and overridable). If the initializer is not found, a nullptr is returned. @param check_outer_scope If true and the graph is a subgraph, check ancestor graph/s for 'name' if not found in 'graph'. @@ -716,7 +716,7 @@ class Graph { return std::find(graph_outputs_.begin(), graph_outputs_.end(), node_arg) != graph_outputs_.end(); } - /** Returns true if one or more of the Node outputs are Graph outputs. + /** Returns true if one or more of the Node outputs are Graph outputs. @remarks Cheaper than calling GetNodeOutputsInGraphOutputs. */ bool NodeProducesGraphOutput(const Node& node) const { @@ -1115,7 +1115,7 @@ class Graph { // Whether to set that no proto sync is required after resolving. // Useful for resolving right after loading from a GraphProto. bool no_proto_sync_required = false; - // When set to true, graph resolve will be called for initialized function bodies as well. This is used + // When set to true, graph resolve will be called for initialized function bodies as well. This is used // in case of nested model local functions. bool traverse_function_body = false; }; @@ -1328,8 +1328,8 @@ class Graph { // so they can be used to resolve outer scope dependencies when running BuildConnections for the subgraphs. common::Status SetOuterScopeNodeArgs(const std::unordered_set& outer_scope_node_args); - // Clear all unused initializers - void CleanUnusedInitializers(const std::unordered_set* initializer_names_to_preserve = nullptr); + // Clear all unused initializers and NodeArgs + void CleanUnusedInitializersAndNodeArgs(const std::unordered_set* initializer_names_to_preserve = nullptr); std::vector CreateNodeArgs(const google::protobuf::RepeatedPtrField& names, const ArgNameToTypeMap& name_to_type_map); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index a6e45ce93a..97be170865 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2678,7 +2678,7 @@ Status Graph::Resolve(const ResolveOptions& options) { // perform the final steps for this graph and all subgraphs auto finalize_func = [&options](Graph& graph) { - graph.CleanUnusedInitializers(options.initializer_names_to_preserve); + graph.CleanUnusedInitializersAndNodeArgs(options.initializer_names_to_preserve); graph.GraphResolveNeeded(false); // if we are resolving immediately after loading from a GraphProto, we don't need to @@ -3345,8 +3345,20 @@ void Graph::ToGraphProtoInternal(ONNX_NAMESPACE::GraphProto& graph_proto) const } } -void Graph::CleanUnusedInitializers(const std::unordered_set* initializer_names_to_preserve) { - std::unordered_set used_args; +void Graph::CleanUnusedInitializersAndNodeArgs(const std::unordered_set* initializer_names_to_preserve) { + // Node Args being used + std::unordered_set used_args; + //Node Args we want to preserved even not being used + std::unordered_set node_args_to_preserve; + if (initializer_names_to_preserve) { + node_args_to_preserve.reserve(initializer_names_to_preserve->size()); + for (const auto& initializer_name : *initializer_names_to_preserve) { + const auto* initializer_node_arg = GetNodeArg(initializer_name); + if (initializer_node_arg != nullptr) { + ORT_IGNORE_RETURN_VALUE(node_args_to_preserve.insert(initializer_node_arg)); + } + } + } // anything that provides a required graph input (GetInputs), an optional graph input (GetOverridableInitializers) // or a graph output (GetOutputs) cannot be removed @@ -3355,35 +3367,36 @@ void Graph::CleanUnusedInitializers(const std::unordered_set* initi const auto& outputs = GetOutputs(); std::for_each(inputs.cbegin(), inputs.cend(), [&used_args](const NodeArg* input) { - ORT_IGNORE_RETURN_VALUE(used_args.insert(input->Name())); + ORT_IGNORE_RETURN_VALUE(used_args.insert(input)); }); std::for_each(overridable_initializers.cbegin(), overridable_initializers.cend(), [&used_args](const NodeArg* input) { - ORT_IGNORE_RETURN_VALUE(used_args.insert(input->Name())); + ORT_IGNORE_RETURN_VALUE(used_args.insert(input)); }); std::for_each(outputs.cbegin(), outputs.cend(), [&used_args](const NodeArg* output) { - ORT_IGNORE_RETURN_VALUE(used_args.insert(output->Name())); + ORT_IGNORE_RETURN_VALUE(used_args.insert(output)); }); for (const auto& node : Nodes()) { for (const auto* def : node.InputDefs()) { - ORT_IGNORE_RETURN_VALUE(used_args.insert(def->Name())); + ORT_IGNORE_RETURN_VALUE(used_args.insert(def)); } for (const auto* def : node.ImplicitInputDefs()) { - ORT_IGNORE_RETURN_VALUE(used_args.insert(def->Name())); + ORT_IGNORE_RETURN_VALUE(used_args.insert(def)); } } std::vector erase_list; - auto end = used_args.end(); + auto used_args_end = used_args.cend(); for (const auto& pv : name_to_initial_tensor_) { const std::string& name = pv.first; - if (used_args.find(name) == end && - (initializer_names_to_preserve == nullptr || - initializer_names_to_preserve->find(name) == initializer_names_to_preserve->cend())) { + const auto* initializer_node_arg = GetNodeArg(name); + ORT_ENFORCE(initializer_node_arg != nullptr, "Cannot find NodeArgs for [", name, "]"); + if (used_args.find(initializer_node_arg) == used_args_end && + node_args_to_preserve.find(initializer_node_arg) == node_args_to_preserve.cend()) { // on the first call to Graph::Resolve we are removing unnecessary initializers that should be removed // from the model. // on later calls we are removing initializers that optimizations have made redundant. @@ -3426,6 +3439,36 @@ void Graph::CleanUnusedInitializers(const std::unordered_set* initi } } }); + + // Clear the unused NodeArgs + // We also want to scan the output NodeArgs of each node + // In case one output of a node is neither used as an input of another node nor an output of graph + for (const auto& node : Nodes()) { + for (const auto* def : node.OutputDefs()) { + ORT_IGNORE_RETURN_VALUE(used_args.insert(def)); + } + } + + // We also need to check the Outer Scope NodeArgs + for (const auto& outer_scope_node_arg_name : outer_scope_node_arg_names_) { + const auto* outer_scope_node_arg = GetNodeArg(outer_scope_node_arg_name); + ORT_ENFORCE(outer_scope_node_arg != nullptr, "Cannot find NodeArgs for [", outer_scope_node_arg_name, "]"); + ORT_IGNORE_RETURN_VALUE(node_args_to_preserve.insert(outer_scope_node_arg)); + } + + auto node_args_to_preserve_end = node_args_to_preserve.cend(); + for (auto it = node_args_.cbegin(), node_args_end = node_args_.cend(); it != node_args_end; /* no increment */) { + auto current_entry = it++; + const auto* current_node_arg = current_entry->second.get(); + const auto& node_arg_name = current_entry->first; + if (!node_arg_name.empty() && used_args.find(current_node_arg) == used_args_end && + node_args_to_preserve.find(current_node_arg) == node_args_to_preserve_end) { + LOGS(logger_, INFO) << "Removing NodeArg '" << node_arg_name << "'. It is no longer used by any node."; + // Need to remove the NodeArg from both value_info_ and node_args_ + value_info_.erase(current_node_arg); + node_args_.erase(current_entry); + } + } } #endif // !defined(ORT_MINIMAL_BUILD) @@ -3823,7 +3866,6 @@ Status Graph::InlineFunction(Node& node) { func_input_output_names.insert(output->Name()); } - // create a uniq_identifier to append to every node name and intermediate input\outputs // to make sure there are no unintended duplicates std::stringstream ss; diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index 8bdce400fd..9403f4ac12 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -873,7 +873,7 @@ TEST_F(GraphTest, GraphConstruction_PriorityBasedTopologicalSort_CompressDecompr node_4 (Identity) decompress (pri = LOCAL_LOW) \ / node_5 (Merge) - | + | */ TypeProto tensor_int32; @@ -943,8 +943,8 @@ TEST_F(GraphTest, GraphConstruction_PriorityBasedTopologicalSort_CompressDecompr \ / | node_8 (Merge) | \ / - node_9 (Merge) - | + node_9 (Merge) + | */ TypeProto tensor_int32; @@ -1012,7 +1012,7 @@ TEST_F(GraphTest, GraphConstruction_PriorityBasedTopologicalSort_Recompute) { node_1 (Identity) recompute_node_1 (pri = LOCAL_LOW) | | node_4 (Identity) | - \ / + \ / node_1_grad (Merge) | */ @@ -1070,7 +1070,7 @@ TEST_F(GraphTest, GraphConstruction_PriorityBasedTopologicalSort_MultiLayerRecom loss (Identity) \ \ \ \ | | \ \ \ 1 | | \ \ - \ / | \ | + \ / | \ | loss_grad recom_node_3 | | \ / | | node_3_grad recom_node_2 | @@ -1256,7 +1256,8 @@ TEST_F(GraphTest, GraphConstruction_CheckGraphInputOutputOrderMaintained) { // Validate that an unused initializer doesn't break graph loading/resolution // and is removed as expected. -TEST_F(GraphTest, UnusedInitializerIsIgnored) { +// Validate unused NodeArgs are removed as expected +TEST_F(GraphTest, UnusedInitializerAndNodeArgsAreIgnored) { Model model("UnusedInitializerIsIgnored", false, *logger_); auto& graph = model.MainGraph(); @@ -1275,17 +1276,28 @@ TEST_F(GraphTest, UnusedInitializerIsIgnored) { graph.AddNode("a", "Identity_Fake", "a", inputs, outputs); TensorProto initializer_tensor; - initializer_tensor.set_name("unused"); + const std::string unused_initializer_name = "unused_initializer"; + initializer_tensor.set_name(unused_initializer_name); initializer_tensor.add_dims(1); initializer_tensor.add_float_data(1.f); initializer_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); graph.AddInitializedTensor(initializer_tensor); ASSERT_TRUE(graph.GetAllInitializedTensors().size() == 1); + ASSERT_NE(nullptr, graph.GetNodeArg(unused_initializer_name)); + + // Add unused NodeArgs + const std::string unused_node_arg_name = graph.GenerateNodeArgName("unused_node_arg"); + ASSERT_EQ(nullptr, graph.GetNodeArg(unused_node_arg_name)); + graph.GetOrCreateNodeArg(unused_node_arg_name, nullptr); + ASSERT_NE(nullptr, graph.GetNodeArg(unused_node_arg_name)); auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()) << status.ErrorMessage(); ASSERT_TRUE(graph.GetAllInitializedTensors().empty()); + ASSERT_EQ(nullptr, graph.GetNodeArg(unused_node_arg_name)); + // Verify NodeArg from the unused initializer is deleted as well + ASSERT_EQ(nullptr, graph.GetNodeArg(unused_initializer_name)); // serialize and reload so we check the loaded from proto path in SetGraphInputsOutputs auto proto = model.ToProto(); @@ -1304,6 +1316,8 @@ TEST_F(GraphTest, UnusedInitializerIsIgnored) { status = graph2.Resolve(); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); ASSERT_TRUE(graph.GetAllInitializedTensors().empty()); + ASSERT_EQ(nullptr, graph.GetNodeArg(unused_node_arg_name)); + ASSERT_EQ(nullptr, graph.GetNodeArg(unused_initializer_name)); } #if !defined(DISABLE_SPARSE_TENSORS)