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
This commit is contained in:
Guoyu Wang 2021-10-01 11:44:26 -07:00 committed by GitHub
parent 8adb9ab85a
commit 60bbdf1403
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 82 additions and 26 deletions

View file

@ -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<std::string>& outer_scope_node_args);
// Clear all unused initializers
void CleanUnusedInitializers(const std::unordered_set<std::string>* initializer_names_to_preserve = nullptr);
// Clear all unused initializers and NodeArgs
void CleanUnusedInitializersAndNodeArgs(const std::unordered_set<std::string>* initializer_names_to_preserve = nullptr);
std::vector<NodeArg*> CreateNodeArgs(const google::protobuf::RepeatedPtrField<std::string>& names,
const ArgNameToTypeMap& name_to_type_map);

View file

@ -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<std::string>* initializer_names_to_preserve) {
std::unordered_set<std::string> used_args;
void Graph::CleanUnusedInitializersAndNodeArgs(const std::unordered_set<std::string>* initializer_names_to_preserve) {
// Node Args being used
std::unordered_set<const NodeArg*> used_args;
//Node Args we want to preserved even not being used
std::unordered_set<const NodeArg*> 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<std::string>* 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<std::string> 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<std::string>* 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;

View file

@ -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)