Update optimizers to be able to utilize a constant initializer from an ancestor graph (#1346)

* Now that we check for a constant initializer in an ancestor graph we also need to be able to retrieve and replace that initializer.
Add helpers to do so.
Update optimizers to use the new helpers.
Fix bug in UnsqueezeElimination where it wasn't checking if the initializer it was replacing was constant.
This commit is contained in:
Scott McKay 2019-07-15 12:41:01 +10:00 committed by GitHub
parent d4ce31ea6d
commit 61b733ce6d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 253 additions and 162 deletions

View file

@ -279,6 +279,10 @@ class Node {
return !attr_to_subgraph_map_.empty();
}
/** Get the const subgraphs from a node.
@remarks Creates a new vector so calling ContainsSubgraphs first is preferred. */
std::vector<gsl::not_null<const Graph*>> GetSubgraphs() const;
/** Gets a map of attribute name to the mutable Graph instances for all subgraphs of the Node.
@returns Map of the attribute name that defines the subgraph to the subgraph's Graph instance.
nullptr if the Node has no subgraphs.
@ -756,6 +760,9 @@ class Graph {
/** Returns the parent graph if this is a subgraph */
const Graph* ParentGraph() const { return parent_graph_; }
/** Returns the mutable parent graph if this is a subgraph */
Graph* MutableParentGraph() { return parent_graph_; }
/** Construct a Graph instance for a subgraph that is created from a GraphProto attribute in a Node.
Inherits some properties from the parent graph.
@param parent_graph The Graph containing the Node which has a GraphProto attribute.

View file

@ -549,6 +549,15 @@ const Graph* Node::GetGraphAttribute(const std::string& attr_name) const {
return const_cast<Node*>(this)->GetMutableGraphAttribute(attr_name);
}
std::vector<gsl::not_null<const Graph*>> Node::GetSubgraphs() const {
std::vector<gsl::not_null<const Graph*>> subgraphs;
subgraphs.reserve(attr_to_subgraph_map_.size());
std::transform(attr_to_subgraph_map_.cbegin(), attr_to_subgraph_map_.cend(), std::back_inserter(subgraphs),
[](const auto& entry) { return entry.second; });
return subgraphs;
}
void Node::ForEachDef(std::function<void(const onnxruntime::NodeArg&, bool is_input)> func,
bool include_missing_optional_defs) const {
for (const auto* arg : InputDefs()) {

View file

@ -54,18 +54,20 @@ static bool OutputEdgeProvidesImplicitInput(const Graph& graph, const GraphEdge&
/** Checks if new_output_name can be used to replace removed_output_name in the subgraph input.
If there is an existing NodeArg in a subgraph that implicitly consumes removed_output_name, it is not safe. */
static bool CanUpdateImplicitInputNameInSubgraph(Node& node,
static bool CanUpdateImplicitInputNameInSubgraph(const Node& node,
const std::string& removed_output_name,
const std::string& new_output_name) {
for (auto& attr_subgraph_pair : node.GetAttributeNameToMutableSubgraphMap()) {
Graph& subgraph = *attr_subgraph_pair.second;
if (!node.ContainsSubgraph())
return true;
for (const gsl::not_null<const Graph*>& subgraph : node.GetSubgraphs()) {
// if we have an existing NodeArg in the subgraph with the new_output_name that would override an implicit input
// with the same name
if (subgraph.GetNodeArg(new_output_name) != nullptr) {
if (subgraph->GetNodeArg(new_output_name) != nullptr) {
return false;
}
for (auto& subgraph_node : attr_subgraph_pair.second->Nodes()) {
for (auto& subgraph_node : subgraph->Nodes()) {
// recurse if this node also consumes removed_output_name as an implicit input (i.e. there are multiple levels of nested
// subgraphs, and at least one level lower uses removed_output_name as an implicit input
const auto subgraph_node_implicit_inputs = subgraph_node.ImplicitInputDefs();
@ -248,6 +250,26 @@ static bool RemoveNodeWithSingleInitializerIn(Graph& graph, Node& node) {
return true;
}
static bool ReplaceInitializerImpl(Graph& graph, const std::string& original_name,
const ONNX_NAMESPACE::TensorProto& initializer, bool check_outer_scope) {
bool replaced = false;
const ONNX_NAMESPACE::TensorProto* old_initializer = nullptr;
if (graph.GetInitializedTensor(original_name, old_initializer)) {
// Be conservative and only remove if the name matches. Graph::CleanupUnusedInitializers can take care
// of removing anything unused after optimization
if (original_name == initializer.name()) {
graph.RemoveInitializedTensor(original_name);
}
graph.AddInitializedTensor(initializer);
replaced = true;
} else if (check_outer_scope && graph.IsSubgraph()) {
replaced = ReplaceInitializerImpl(*graph.MutableParentGraph(), original_name, initializer, check_outer_scope);
}
return replaced;
}
//----------------------------
//--- end of local helpers ---
//----------------------------
@ -292,52 +314,6 @@ bool IsSupportedProvider(const Node& node,
compatible_providers.find(node.GetExecutionProviderType()) == compatible_providers.end());
}
Status ForAllMutableSubgraphs(Graph& graph, std::function<Status(Graph&)> func) {
Status status = Status::OK();
for (auto& node : graph.Nodes()) {
for (auto& attr_name_to_subgraph_pair : node.GetAttributeNameToMutableSubgraphMap()) {
Graph* subgraph = attr_name_to_subgraph_pair.second;
ORT_ENFORCE(subgraph, "Main Graph instance should have populated all subgraphs when being resolved.");
status = func(*subgraph);
ORT_RETURN_IF_ERROR(status);
// recurse
status = ForAllMutableSubgraphs(*subgraph, func);
ORT_RETURN_IF_ERROR(status);
}
}
return status;
}
Status ForAllSubgraphs(const Graph& graph, std::function<Status(const Graph&)> func) {
Status status = Status::OK();
for (auto& node : graph.Nodes()) {
for (auto& attribute : node.GetAttributes()) {
auto& name = attribute.first;
auto& proto = attribute.second;
// check if it has a subgraph
if (proto.has_g()) {
const Graph* subgraph = node.GetGraphAttribute(name);
ORT_ENFORCE(subgraph, "Main Graph instance should have populated all subgraphs when being resolved.");
status = func(*subgraph);
ORT_RETURN_IF_ERROR(status);
// recurse
status = ForAllSubgraphs(*subgraph, func);
ORT_RETURN_IF_ERROR(status);
}
}
}
return status;
}
bool IsSingleInSingleOutNode(const Node& node) {
return node.InputDefs().size() == 1 && node.ImplicitInputDefs().empty() && node.OutputDefs().size() == 1;
}
@ -399,48 +375,69 @@ bool IsGraphInput(const Graph& graph, const NodeArg* input) {
return std::find(graph_inputs.begin(), graph_inputs.end(), input) != graph_inputs.end();
}
bool IsConstantInitializer(const Graph& graph, const std::string& initializer_name, bool check_outer_scope) {
const onnx::TensorProto* initializer = nullptr;
bool is_local_initializer = graph.GetInitializedTensor(initializer_name, initializer);
// if we know it's an initializer we assume it's constant initially.
// otherwise it's not an initializer so can't be a constant initializer by definition.
bool constant_initializer = is_local_initializer;
if (is_local_initializer) {
const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const Graph& graph, const std::string& initializer_name,
bool check_outer_scope) {
const ONNX_NAMESPACE::TensorProto* initializer = nullptr;
if (graph.GetInitializedTensor(initializer_name, initializer)) {
if (graph.CanOverrideInitializer()) {
const auto& graph_inputs = graph.GetInputsIncludingInitializers();
constant_initializer = std::none_of(graph_inputs.cbegin(), graph_inputs.cend(),
[&initializer_name](const NodeArg* input) {
return input->Name() == initializer_name;
});
bool is_constant = std::none_of(graph_inputs.cbegin(), graph_inputs.cend(),
[&initializer_name](const NodeArg* input) {
return input->Name() == initializer_name;
});
if (!is_constant) {
initializer = nullptr;
}
}
} else if (check_outer_scope && graph.IsSubgraph()) {
constant_initializer = IsConstantInitializer(*graph.ParentGraph(), initializer_name, check_outer_scope);
initializer = GetConstantInitializer(*graph.ParentGraph(), initializer_name);
}
return constant_initializer;
return initializer;
}
bool IsConstantInitializer(const Graph& graph, const std::string& initializer_name, bool check_outer_scope) {
const onnx::TensorProto* initializer = GetConstantInitializer(graph, initializer_name, check_outer_scope);
return initializer != nullptr;
}
bool NodeArgIsConstant(const Graph& graph, const NodeArg& node_arg) {
return IsConstantInitializer(graph, node_arg.Name(), true);
}
bool AllNodeInputsAreConstant(const Graph& graph, const Node& node) {
bool AllNodeInputsAreConstant(const Graph& graph, const Node& node, InitializedTensorSet& constant_inputs) {
// clear so we have a known state. if we fail part way through we go back to this state.
constant_inputs.clear();
// only initializers can be constant, and there's no edge from a node to an initializer
// so the input edges count must be 0
if (node.GetInputEdgesCount() > 0) {
return false;
}
for (const auto* input_def : node.InputDefs()) {
// Important note: when an initializer appears in the graph's input, this input will not be considered constant,
// because it can be overriden by the user at runtime. For constant folding to be applied, the initializer should
// because it can be overridden by the user at runtime. For constant folding to be applied, the initializer should
// not appear in the graph's inputs (that is the only way to guarantee it will always be constant).
if (!NodeArgIsConstant(graph, *input_def)) {
const ONNX_NAMESPACE::TensorProto* initializer = GetConstantInitializer(graph, input_def->Name(), true);
if (initializer) {
constant_inputs.insert({input_def->Name(), initializer});
} else {
constant_inputs.clear();
return false;
}
}
return true;
}
void ReplaceInitializer(Graph& graph, const std::string& original_name, const ONNX_NAMESPACE::TensorProto& initializer,
bool check_outer_scope) {
ORT_ENFORCE(ReplaceInitializerImpl(graph, original_name, initializer, check_outer_scope),
"Failed to replace initializer. Original initializer was not found. Name:", original_name);
}
size_t RemoveNodeOutputEdges(Graph& graph, Node& node) {
std::vector<GraphEdge> output_edges = GetNodeOutputEdges(node);
RemoveGraphEdges(graph, output_edges);

View file

@ -39,16 +39,31 @@ bool IsOutputUsed(const Node& node, int index);
bool IsGraphInput(const Graph& graph, const NodeArg* input);
/** returns true if 'name' is an initializer, and is constant and cannot be overridden at runtime.
@param check_outer_scope If true and the graph is a subgraph, check parent graph/s for 'name' if not found in 'graph'.
@param check_outer_scope If true and the graph is a subgraph, check ancestor graph/s for 'name' if not found in 'graph'.
*/
bool IsConstantInitializer(const Graph& graph, const std::string& name, bool check_outer_scope);
bool IsConstantInitializer(const Graph& graph, const std::string& name, bool check_outer_scope = true);
/** Checks if the given node has only constant inputs (initializers). */
bool AllNodeInputsAreConstant(const Graph& graph, const Node& node);
/** returns the initializer's TensorProto if 'name' is an initializer, and is constant and
cannot be overridden at runtime. If the initializer is not found or is not constant 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'.
*/
const ONNX_NAMESPACE::TensorProto* GetConstantInitializer(const Graph& graph, const std::string& name,
bool check_outer_scope = true);
/** Find the initializer called 'original_name' in 'graph', or its ancestors if check_outer_scope is true,
and replace with 'initializer' in the current graph.
Does NOT look in any subgraphs. Requires original_name to match an initializer.
*/
void ReplaceInitializer(Graph& graph, const std::string& original_name, const ONNX_NAMESPACE::TensorProto& initializer,
bool check_outer_scope = true);
/** Checks if the given NodeArg is constant, i.e., it appears in the graph's initializers but not in its inputs. */
bool NodeArgIsConstant(const Graph& graph, const NodeArg& node_arg);
/** Checks if the given node has only constant inputs (initializers) and if so returns them in constant_inputs as they
may come from outer scope. */
bool AllNodeInputsAreConstant(const Graph& graph, const Node& node, InitializedTensorSet& constant_inputs);
/** Gets the name of the incoming NodeArg with the specified index for the given node. */
const std::string& GetNodeInputName(const Node& node, int index);
@ -71,9 +86,6 @@ bool GetRepeatedNodeAttributeValues(const Node& node,
return false;
}
Status ForAllMutableSubgraphs(Graph& main_graph, std::function<Status(Graph&)> func);
Status ForAllSubgraphs(const Graph& main_graph, std::function<Status(const Graph&)> func);
/** Removes the given Node from the Graph and keeps Graph consistent by rebuilding needed connections.
We support the removal of the Node as long as the following conditions hold:
- There should be no implicit inputs.

View file

@ -23,23 +23,24 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level)
ORT_RETURN_IF_ERROR(Recurse(*node, modified, graph_level));
InitializedTensorSet constant_inputs;
// Check if constant folding can be applied on this node.
if (!graph_utils::IsSupportedProvider(*node, GetCompatibleExecutionProviders()) ||
excluded_op_types_.find(node->OpType()) != excluded_op_types_.end() ||
// constant folding is not currently supported for nodes that include subgraphs (control flow operators,
// such as If/Loop/Scan, fall into this category).
// constant folding does not support executing a node that includes subgraphs (control flow operators,
// such as If/Loop/Scan, fall into this category). individual nodes in the subgraph will be processed
// by the Recurse call above
node->ContainsSubgraph() ||
// if the node output is in the graph output, we will get a graph with no nodes.
// TODO check if this is allowed in ONNX and ORT.
graph.IsNodeOutputsInGraphOutputs(*node) ||
!graph_utils::AllNodeInputsAreConstant(graph, *node)) {
!graph_utils::AllNodeInputsAreConstant(graph, *node, constant_inputs)) {
continue;
}
// Create execution frame for executing constant nodes.
// NOTE: As we call AllNodeInputsAreConstant we can use the full list of initializers from
// graph.GetAllInitializedTensors() without filtering out overridable (i.e. non-constant) initializers
OptimizerExecutionFrame::Info info({node}, graph.GetAllInitializedTensors());
OptimizerExecutionFrame::Info info({node}, constant_inputs);
std::vector<int> fetch_mlvalue_idxs;
for (const auto* node_out : node->OutputDefs()) {

View file

@ -15,19 +15,14 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie
const auto& conv_inputs = conv_node.InputDefs();
const auto& add_inputs = add_node.InputDefs();
const ONNX_NAMESPACE::TensorProto* conv_W_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(conv_inputs[1]->Name(), conv_W_tensor_proto)) {
return Status::OK();
}
const auto* conv_W_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[1]->Name());
ORT_ENFORCE(conv_W_tensor_proto);
const ONNX_NAMESPACE::TensorProto* add_B_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(add_inputs[1]->Name(), add_B_tensor_proto)) {
return Status::OK();
}
const auto* add_B_tensor_proto = graph_utils::GetConstantInitializer(graph, add_inputs[1]->Name());
ORT_ENFORCE(add_B_tensor_proto);
// Currently, fusion is only supported for float or double data type.
if (!Initializer::IsSupportedDataType(add_B_tensor_proto) ||
conv_W_tensor_proto->dims_size() < 4) {
if (!Initializer::IsSupportedDataType(add_B_tensor_proto) || conv_W_tensor_proto->dims_size() < 4) {
return Status::OK();
}
@ -51,11 +46,9 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie
}
}
const ONNX_NAMESPACE::TensorProto* conv_B_tensor_proto = nullptr;
if (conv_inputs.size() == 3) {
if (!graph.GetInitializedTensor(conv_inputs[2]->Name(), conv_B_tensor_proto)) {
return Status::OK();
}
const auto* conv_B_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[2]->Name());
ORT_ENFORCE(conv_B_tensor_proto);
if (!Initializer::IsSupportedDataType(conv_B_tensor_proto) ||
conv_B_tensor_proto->data_type() != add_B_tensor_proto->data_type() ||
@ -78,8 +71,7 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie
conv_B->ToProto(&new_conv_B_tensor_proto);
// Replace initializers of conv node
graph.RemoveInitializedTensor(conv_inputs[2]->Name());
graph.AddInitializedTensor(new_conv_B_tensor_proto);
graph_utils::ReplaceInitializer(graph, conv_inputs[2]->Name(), new_conv_B_tensor_proto);
} else {
NodeArg* add_B_node_arg = graph.GetNodeArg(add_B_tensor_proto->name());
if (add_B_node_arg == nullptr) {
@ -92,8 +84,7 @@ Status ConvAddFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& modifie
new_conv_B_tensor_proto.clear_dims();
new_conv_B_tensor_proto.add_dims(dim);
graph.RemoveInitializedTensor(add_B_tensor_proto->name());
graph.AddInitializedTensor(new_conv_B_tensor_proto);
graph_utils::ReplaceInitializer(graph, add_B_tensor_proto->name(), new_conv_B_tensor_proto);
// Update shape of NodeArg
TensorShapeProto shape;

View file

@ -23,31 +23,21 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff
// Get initializers of BatchNormalization
const auto& bn_inputs = bn_node.InputDefs();
const ONNX_NAMESPACE::TensorProto* bn_scale_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(bn_inputs[1]->Name(), bn_scale_tensor_proto)) {
return Status::OK();
}
const auto* bn_scale_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[1]->Name());
ORT_ENFORCE(bn_scale_tensor_proto);
const ONNX_NAMESPACE::TensorProto* bn_B_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(bn_inputs[2]->Name(), bn_B_tensor_proto)) {
return Status::OK();
}
const auto* bn_B_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[2]->Name());
ORT_ENFORCE(bn_B_tensor_proto);
const ONNX_NAMESPACE::TensorProto* bn_mean_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(bn_inputs[3]->Name(), bn_mean_tensor_proto)) {
return Status::OK();
}
const auto* bn_mean_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[3]->Name());
ORT_ENFORCE(bn_mean_tensor_proto);
const ONNX_NAMESPACE::TensorProto* bn_var_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(bn_inputs[4]->Name(), bn_var_tensor_proto)) {
return Status::OK();
}
const auto* bn_var_tensor_proto = graph_utils::GetConstantInitializer(graph, bn_inputs[4]->Name());
ORT_ENFORCE(bn_var_tensor_proto);
const auto& conv_inputs = conv_node.InputDefs();
const ONNX_NAMESPACE::TensorProto* conv_W_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(conv_inputs[1]->Name(), conv_W_tensor_proto)) {
return Status::OK();
}
const auto* conv_W_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[1]->Name());
ORT_ENFORCE(conv_W_tensor_proto);
// Currently, fusion is only supported for float or double data type.
if (!Initializer::IsSupportedDataType(bn_scale_tensor_proto) ||
@ -76,12 +66,11 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff
auto bn_var = std::make_unique<Initializer>(bn_var_tensor_proto);
auto conv_W = std::make_unique<Initializer>(conv_W_tensor_proto);
const ONNX_NAMESPACE::TensorProto* conv_B_tensor_proto = nullptr;
std::unique_ptr<Initializer> conv_B = nullptr;
const ONNX_NAMESPACE::TensorProto* conv_B_tensor_proto = nullptr;
if (conv_inputs.size() == 3) {
if (!graph.GetInitializedTensor(conv_inputs[2]->Name(), conv_B_tensor_proto)) {
return Status::OK();
}
conv_B_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[2]->Name());
ORT_ENFORCE(conv_B_tensor_proto);
if (!Initializer::IsSupportedDataType(conv_B_tensor_proto) ||
conv_B_tensor_proto->dims_size() != 1 ||
@ -124,24 +113,23 @@ Status ConvBNFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_eff
}
// Replace initializers of conv node
graph.RemoveInitializedTensor(conv_W_tensor_proto->name());
graph_utils::ReplaceInitializer(graph, conv_W_tensor_proto->name(), new_conv_W_tensor_proto);
if (conv_inputs.size() == 3) {
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 6011) // Not deferencing null pointer. conv_B_tensor_proto is set on line 93
#endif
graph.RemoveInitializedTensor(conv_B_tensor_proto->name());
graph_utils::ReplaceInitializer(graph, conv_B_tensor_proto->name(), new_conv_B_tensor_proto);
#ifdef _MSC_VER
#pragma warning(pop)
#endif
} else {
graph.RemoveInitializedTensor(bn_B_tensor_proto->name());
graph_utils::ReplaceInitializer(graph, bn_B_tensor_proto->name(), new_conv_B_tensor_proto);
conv_node.MutableInputDefs().push_back(bn_B_node_arg);
conv_node.MutableInputArgsCount()[2] = 1;
}
graph.AddInitializedTensor(new_conv_W_tensor_proto);
graph.AddInitializedTensor(new_conv_B_tensor_proto);
// Remove BN node.
auto* bn_node_to_remove = graph.GetNode(bn_node.Index());

View file

@ -15,15 +15,11 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef
const auto& conv_inputs = conv_node.InputDefs();
const auto& mul_inputs = mul_node.InputDefs();
const ONNX_NAMESPACE::TensorProto* conv_W_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(conv_inputs[1]->Name(), conv_W_tensor_proto)) {
return Status::OK();
}
const auto* conv_W_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[1]->Name());
ORT_ENFORCE(conv_W_tensor_proto);
const ONNX_NAMESPACE::TensorProto* mul_B_tensor_proto = nullptr;
if (!graph.GetInitializedTensor(mul_inputs[1]->Name(), mul_B_tensor_proto)) {
return Status::OK();
}
const auto* mul_B_tensor_proto = graph_utils::GetConstantInitializer(graph, mul_inputs[1]->Name());
ORT_ENFORCE(mul_B_tensor_proto);
if (!Initializer::IsSupportedDataType(conv_W_tensor_proto) ||
!Initializer::IsSupportedDataType(mul_B_tensor_proto) ||
@ -61,10 +57,9 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef
std::unique_ptr<Initializer> conv_B = nullptr;
const bool is_3d = conv_inputs.size() == 3;
if (is_3d) {
if (!graph.GetInitializedTensor(conv_inputs[2]->Name(), conv_B_tensor_proto))
return Status::OK();
if (conv_B_tensor_proto == nullptr)
return Status(ONNXRUNTIME, FAIL, "Internal error in ConvMulFusion. conv_B_tensor_proto is NULL");
conv_B_tensor_proto = graph_utils::GetConstantInitializer(graph, conv_inputs[2]->Name());
ORT_ENFORCE(conv_B_tensor_proto);
if (!Initializer::IsSupportedDataType(conv_B_tensor_proto) ||
conv_B_tensor_proto->data_type() != mul_B_tensor_proto->data_type() ||
conv_B_tensor_proto->dims_size() != 1 ||
@ -90,14 +85,12 @@ Status ConvMulFusion::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_ef
conv_W->ToProto(&new_conv_W_tensor_proto);
// Replace initializers of conv node
graph.RemoveInitializedTensor(conv_inputs[1]->Name());
graph.AddInitializedTensor(new_conv_W_tensor_proto);
graph_utils::ReplaceInitializer(graph, conv_inputs[1]->Name(), new_conv_W_tensor_proto);
if (is_3d) {
ONNX_NAMESPACE::TensorProto new_conv_B_tensor_proto(*conv_B_tensor_proto);
conv_B->ToProto(&new_conv_B_tensor_proto);
graph.RemoveInitializedTensor(conv_inputs[2]->Name());
graph.AddInitializedTensor(new_conv_B_tensor_proto);
graph_utils::ReplaceInitializer(graph, conv_inputs[2]->Name(), new_conv_B_tensor_proto);
}
// Remove Mul node.

View file

@ -25,11 +25,9 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect&
// Generate new dims.
NodeArg* input_def = node.MutableInputDefs()[0];
const ONNX_NAMESPACE::TensorProto* tensor_proto = nullptr;
graph.GetInitializedTensor(input_def->Name(), tensor_proto);
if (tensor_proto == nullptr) {
return Status::OK();
}
const auto* tensor_proto = graph_utils::GetConstantInitializer(graph, input_def->Name());
ORT_ENFORCE(tensor_proto);
std::vector<int64_t> new_dims(axes.size() + tensor_proto->dims().size(), 0);
if (new_dims.size() >= std::numeric_limits<int>::max()) {
return Status(ONNXRUNTIME, FAIL, "index out of range");
@ -56,8 +54,13 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect&
new_tensor_proto.add_dims(new_dims[i]);
}
}
graph.RemoveInitializedTensor(input_def->Name());
graph.AddInitializedTensor(new_tensor_proto);
// TODO: This seems wrong as there's no check whether another node is using the initializer before replacing it.
// Shouldn't we check that or alternatively create an initializer with a different name and let
// Graph::CleanUnusedInitializers remove the original one if nothing else consumes it?
// graph.RemoveInitializedTensor(input_def->Name());
// graph.AddInitializedTensor(new_tensor_proto);
graph_utils::ReplaceInitializer(graph, input_def->Name(), new_tensor_proto);
// Update shape of NodeArg.
TensorShapeProto shape;
@ -75,8 +78,8 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect&
} // namespace onnxruntime
bool UnsqueezeElimination::SatisfyCondition(const Graph& graph, const Node& node) const {
// Attempt to remove an Unsqueeze operator only if it gets an initializer as input.
return node.GetInputEdgesCount() == 0 &&
// Attempt to remove an Unsqueeze operator only if it gets a constant initializer as input.
return graph_utils::IsConstantInitializer(graph, node.InputDefs()[0]->Name()) &&
!graph.IsNodeOutputsInGraphOutputs(node);
}

View file

@ -125,7 +125,7 @@ TEST(SessionStateTest, TestInitializerProcessing) {
bool found = initialized_tensors.find(idx) != initialized_tensors.cend();
ASSERT_TRUE(found) << "Missing entry for " << entry.first << " in session state initialized tensors";
if (graph_utils::IsConstantInitializer(graph, entry.first, true)) {
if (graph_utils::IsConstantInitializer(graph, entry.first, false)) {
found = const_initialized_tensors.find(idx) != const_initialized_tensors.cend();
ASSERT_TRUE(found) << "Missing entry for " << entry.first << " in session state const initialized tensors";
}

View file

@ -42,15 +42,30 @@ IExecutionProvider* TestNnapiExecutionProvider() {
}
#endif
static void CountOpsInGraphImpl(const Graph& graph, std::map<std::string, int>& ops) {
for (auto& node : graph.Nodes()) {
auto pos = ops.find(node.OpType());
if (pos == ops.end()) {
ops[node.OpType()] = 1;
} else {
++pos->second;
}
if (node.ContainsSubgraph()) {
for (auto& subgraph : node.GetSubgraphs()) {
CountOpsInGraphImpl(*subgraph, ops);
}
}
}
}
// Returns a map with the number of occurrences of each operator in the graph.
// Helper function to check that the graph transformations have been successfully applied.
std::map<std::string, int> CountOpsInGraph(const Graph& graph) {
std::map<std::string, int> op_to_count;
for (auto& node : graph.Nodes()) {
op_to_count[node.OpType()] =
op_to_count.count(node.OpType()) == 0 ? 1 : ++op_to_count[node.OpType()];
}
return op_to_count;
std::map<std::string, int> ops;
CountOpsInGraphImpl(graph, ops);
return ops;
}
} // namespace test

View file

@ -115,6 +115,80 @@ TEST(GraphTransformationTests, ConstantFolding) {
ASSERT_TRUE(op_to_count["Unsqueeze"] == 0);
}
TEST(GraphTransformationTests, ConstantFoldingSubgraph) {
TensorProto value_tensor;
value_tensor.add_dims(1);
value_tensor.add_float_data(1.f);
value_tensor.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
TypeProto float_tensor_type;
float_tensor_type.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
float_tensor_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1);
auto create_subgraph = [&](GraphProto& graph_proto) {
// create subgraph that has an Add node to add a local and parent graph initializer
Model model("ConstantFoldingSubgraphTest_subgraph");
auto& graph = model.MainGraph();
TensorProto local_constant(value_tensor);
local_constant.set_name("local_constant");
graph.AddInitializedTensor(local_constant);
auto& local_constant_arg = graph.GetOrCreateNodeArg("local_constant", &float_tensor_type);
auto& parent_constant_arg = graph.GetOrCreateNodeArg("parent_constant", &float_tensor_type);
graph.AddOuterScopeNodeArg("parent_constant");
auto& add_out = graph.GetOrCreateNodeArg("add_out", &float_tensor_type);
graph.AddNode("add", "Add", "Add two inputs.", {&parent_constant_arg, &local_constant_arg}, {&add_out});
auto& subgraph_out = graph.GetOrCreateNodeArg("subgraph_out", &float_tensor_type);
graph.AddNode("identity", "Identity", "So Add isn't providing graph output.", {&add_out}, {&subgraph_out});
auto status = graph.Resolve();
ASSERT_TRUE(status.IsOK()) << status;
graph_proto = graph.ToGraphProto();
};
Model model("ConstantFoldingSubgraphTest_main_graph");
auto& graph = model.MainGraph();
// add initializer at parent level
TensorProto parent_value_tensor(value_tensor);
parent_value_tensor.set_name("parent_constant");
graph.AddInitializedTensor(parent_value_tensor);
// put the subgraph in an If node
TypeProto if_cond_type;
if_cond_type.mutable_tensor_type()->set_elem_type(TensorProto_DataType_BOOL);
if_cond_type.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1);
auto& if_cond_input = graph.GetOrCreateNodeArg("if_in", &if_cond_type);
auto& if_output = graph.GetOrCreateNodeArg("if_out", &float_tensor_type);
auto& if_node = graph.AddNode("if", "If", "If node", {&if_cond_input}, {&if_output});
GraphProto subgraph;
create_subgraph(subgraph);
if_node.AddAttribute("then_branch", {subgraph});
if_node.AddAttribute("else_branch", {subgraph});
auto status = graph.Resolve();
ASSERT_TRUE(status.IsOK()) << status;
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Add"] == 2); // one in each subgraph
onnxruntime::GraphTransformerManager graph_transformation_mgr{5};
graph_transformation_mgr.Register(std::make_unique<ConstantFolding>(), TransformerLevel::Level1);
status = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level1);
ASSERT_TRUE(status.IsOK()) << status;
op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Add"] == 0)
<< "Constant folding should have been able to remove the Add node in both subgraphs";
}
TEST(GraphTransformationTests, ShapeToInitializer) {
string model_uri = MODEL_FOLDER + "shape-add.onnx";
std::shared_ptr<Model> model;
@ -308,9 +382,10 @@ TEST(GraphTransformationTests, NegativeFuseConvAddNoBias) {
ASSERT_TRUE(graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2).IsOK());
// Nodes are not fused because the weights to conv/add are not constants (they appear in the graph inputs).
// Unsqueeze is also not eliminated as the initializer that is its input is also not constant
std::map<std::string, int> op_to_count = CountOpsInGraph(graph);
ASSERT_TRUE(op_to_count["Add"] != 0);
ASSERT_TRUE(op_to_count["Unsqueeze"] == 0);
ASSERT_TRUE(op_to_count["Unsqueeze"] != 0);
}
TEST(GraphTransformationTests, FuseConvAddMul3D) {