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

Replace ORT's function shape inference with ONNX's (#11538)

Browse source 
* Function inlining tests

* Replace ORT copy of function shape inference

* Remove std::move

* Fix memory leak

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>

* Address feedback
This commit is contained in:
G. Ramalingam 2022-06-03 12:54:28 -07:00 committed by GitHub
parent ec05313cd9
commit 98960c53fe
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 138 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

136
onnxruntime/core/graph/function_utils.cc
View file

@ -1,5 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <unordered_map>
#include "core/graph/function_utils.h"
#include "core/common/inlined_containers.h"
#include "core/framework/tensorprotoutils.h"
@ -7,131 +9,6 @@
#include "core/graph/function_impl.h"
#include "core/graph/model_load_utils.h"
namespace ONNX_NAMESPACE {
// Infer shape for functions. This also supports
// nested model local functions.
// TODO: Add this to onnx instead of adding it here.
void InferShapeForFunctionNode(
InferenceContext& ctx,
const FunctionProto& func_proto,
const onnxruntime::InlinedHashMap<std::string, int>& func_opset_imports,
const ShapeInferenceOptions& options,
const ISchemaRegistry* schema_registry,
const onnxruntime::InlinedHashMap<std::string, const FunctionProto*>& in_model_functions,
std::function<std::string(const std::string& function_domain, const std::string& function_name)> get_func_id) {
GraphProto g;
// Get a temporary tensor-shape map
const auto num_func_inputs = func_proto.input_size();
std::unordered_map<std::string, TypeProto*> value_types_by_name;
onnxruntime::InlinedVector<TypeProto> types_cache;
types_cache.reserve(func_proto.input_size());
for (int i = 0; i < num_func_inputs; ++i) {
types_cache.push_back(*ctx.getInputType(i));
value_types_by_name[func_proto.input().Get(i)] = &types_cache[i];
}
// Get a temporary initial value map
std::unordered_map<std::string, const TensorProto*> initializers_by_name;
std::unordered_map<std::string, const SparseTensorProto*> sparse_initializers_by_name;
for (int i = 0; i < static_cast<int>(ctx.getNumInputs()) && i < num_func_inputs; ++i) {
const TypeProto* type = ctx.getInputType(i);
if (type->value_case() == TypeProto::kTensorType && ctx.getInputData(i) != nullptr) {
initializers_by_name[func_proto.input().Get(i)] = ctx.getInputData(i);
} else if (type->value_case() == TypeProto::kSparseTensorType &&
ctx.getInputSparseData(i) != nullptr) {
sparse_initializers_by_name[func_proto.input().Get(i)] = ctx.getInputSparseData(i);
}
}
std::unordered_map<std::string, const AttributeProto*> attr_map;
for (auto& attr : func_proto.attribute()) {
if (ctx.getAttribute(attr) != nullptr) {
attr_map[attr] = ctx.getAttribute(attr);
}
}
for (auto& n : func_proto.node()) {
NodeProto copy_n(n);
// Add attribute information into the temporary node
copy_n.clear_attribute();
for (const auto& attr : n.attribute()) {
if (attr.has_ref_attr_name()) {
if (attr_map.count(attr.ref_attr_name())) {
auto copy_attr = *attr_map[attr.ref_attr_name()];
copy_attr.set_name(attr.name());
copy_n.add_attribute()->CopyFrom(copy_attr);
}
} else {
copy_n.add_attribute()->CopyFrom(attr);
}
}
ONNX_NAMESPACE::shape_inference::InferenceContextImpl func_node_ctx(
copy_n, value_types_by_name, initializers_by_name, sparse_initializers_by_name, {});
// Resolve domain for node
auto it = func_opset_imports.find(n.domain());
if (it == func_opset_imports.end()) {
fail_type_inference("Cannot infer type and shape for function", func_proto.name(),
". No opset import for domain", n.domain(), " referenced by function body node ",
n.name(), " optype ", n.op_type());
}
auto domain_version = it->second;
const auto schema = schema_registry->GetSchema(n.op_type(), domain_version, n.domain());
if (schema) {
schema->GetTypeAndShapeInferenceFunction()(func_node_ctx);
} else {
// check model local functions for FunctionProto
auto iter = in_model_functions.find(get_func_id(n.domain(), n.op_type()));
if (iter == in_model_functions.end()) {
return;
}
onnxruntime::InlinedHashMap<std::string, int> func_node_opset_imports;
for (const auto& opset_import : iter->second->opset_import()) {
// If graph imports does not contain opset_import then insert it otherwise the one in graph imports overrides.
// If the opset imports are not compatible then this will be caught during function body inline.
func_node_opset_imports.insert({opset_import.domain(), static_cast<int>(opset_import.version())});
}
InferShapeForFunctionNode(func_node_ctx, *iter->second, func_node_opset_imports, options, schema_registry, in_model_functions, get_func_id);
}
for (int i = 0; i < copy_n.output_size(); ++i) {
TypeProto* inferred_output_type = func_node_ctx.getOutputType(i);
// Checking, Storing the inferred information
auto iter = value_types_by_name.find(n.output(i));
TypeProto* existingType = nullptr;
if (iter != value_types_by_name.end()) {
existingType = iter->second;
shape_inference::checkShapesAndTypes(*inferred_output_type, *existingType);
} else {
// Store the inferred type info in the
// subgraph temporarily
auto vi = g.add_value_info();
vi->set_name(copy_n.output(i));
existingType = vi->mutable_type();
}
shape_inference::mergeShapesAndTypes(*inferred_output_type, existingType);
// Make merged info available to further inference.
value_types_by_name[copy_n.output(i)] = existingType;
}
}
for (int i = 0; i < func_proto.output_size(); ++i) {
const std::string& output_name = func_proto.output().Get(i);
// Skip if no type inferred for the tensor
auto iter = value_types_by_name.find(output_name);
if (iter != value_types_by_name.cend()) {
// Copy the type info to ctx
// to pass back to main graph
auto type_proto = ctx.getOutputType(i);
type_proto->CopyFrom(*(iter->second));
}
}
}
} // namespace ONNX_NAMESPACE
namespace onnxruntime {
namespace function_utils {
@ -415,7 +292,7 @@ std::unique_ptr<ONNX_NAMESPACE::OpSchema> CreateSchema(const std::string& functi
const auto onnx_released_versions =
schema_registry.GetLastReleasedOpsetVersions(false);
InlinedHashMap<std::string, int> func_domain_to_version;
std::unordered_map<std::string, int> func_domain_to_version;
for (auto& opSet : onnx_func_proto->opset_import()) {
const auto& domain = opSet.domain();
const auto version = gsl::narrow_cast<int>(opSet.version());
@ -437,7 +314,10 @@ std::unique_ptr<ONNX_NAMESPACE::OpSchema> CreateSchema(const std::string& functi
[onnx_func_proto, func_domain_to_version, &model_local_functions](ONNX_NAMESPACE::InferenceContext& ctx) {
auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
ONNX_NAMESPACE::ShapeInferenceOptions options{true, 1, false};
InferShapeForFunctionNode(ctx, *onnx_func_proto, func_domain_to_version, options, schema_registry, model_local_functions, function_utils::GetFunctionIdentifier);
std::unordered_map<std::string, const ONNX_NAMESPACE::FunctionProto*> map_copy(model_local_functions.begin(),
model_local_functions.end());
ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(*onnx_func_proto, func_domain_to_version,
schema_registry, ctx, options, map_copy);
});
op_schema->Finalize();
@ -661,4 +541,4 @@ Status Instantiate(onnxruntime::Graph& graph,
}
}
}
}

130
onnxruntime/test/framework/function_test.cc Normal file
View file

@ -0,0 +1,130 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "gtest/gtest.h"
#include "onnx/defs/parser.h"
#include "core/graph/model.h"
#include "core/providers/cpu/cpu_execution_provider.h"
#include "core/session/inference_session.h"
#include "test/test_environment.h"
#include "test/framework/test_utils.h"
#include "test/common/tensor_op_test_utils.h"
// Unit tests to check the implementation of functions, model-local functions,
// function-inlining etc.
namespace onnxruntime {
namespace test {
static void Check(const char* source,
const char* input_name, std::vector<float> input_values,
const char* output_name, std::vector<float> output_values) {
// Convert source-representation of model to ModelProto:
ONNX_NAMESPACE::OnnxParser parser(source);
ONNX_NAMESPACE::ModelProto model;
auto parse_status = parser.Parse(model);
EXPECT_TRUE(parse_status.IsOK()) << parse_status.ErrorMessage();
EXPECT_TRUE(parser.EndOfInput()) << "Extra unparsed input unexpected.";
// Serialize and then load model:
std::string serialized_model;
const bool serialization_status = model.SerializeToString(&serialized_model);
EXPECT_TRUE(serialization_status) << "Failed to serialize proto to string";
SessionOptions session_options;
InferenceSession session_object{session_options, GetEnvironment()};
std::stringstream sstr(serialized_model);
auto status = session_object.Load(sstr);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
status = session_object.Initialize();
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
RunOptions run_options;
run_options.run_tag = session_options.session_logid;
NameMLValMap feeds;
std::unique_ptr<CPUExecutionProvider> provider = std::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo());
OrtValue ort_value;
CreateMLValue<float>(provider->GetAllocator(0, OrtMemTypeDefault), {int64_t(input_values.size())}, input_values, &ort_value);
feeds.insert(std::make_pair(std::string(input_name), ort_value));
std::vector<OrtValue> fetches;
status = session_object.Run(run_options, feeds, {output_name}, &fetches);
EXPECT_TRUE(status.IsOK()) << "Session Run failed.";
auto& tensor = fetches[0].Get<Tensor>();
size_t size = static_cast<size_t>(tensor.Shape().Size());
EXPECT_EQ(size, output_values.size());
auto* data = tensor.template Data<float>();
float threshold = 0.001f;
for (size_t i = 0; i < size; ++i) {
ASSERT_NEAR(data[i], output_values[i], threshold) << "at position i:" << i;
}
}
TEST(FunctionTest, Basic) {
const char* code = R"(
<
ir_version: 8,
opset_import: [ "" : 16, "local" : 1 ]
>
agraph (float[N] x) => (float[N] y)
{
y = local.myfun (x)
}
<
opset_import: [ "" : 16 ],
domain: "local"
>
myfun (lx) => (ly) {
two = Constant <value = float[1] {2.0}> ()
ly = Mul (lx, two)
}
)";
Check(code, "x", {1.0, 2.0, 3.0}, "y", {2.0, 4.0, 6.0});
}
// Check that variables are renamed to avoid conflicts when multiple
// calls are inlined.
TEST(FunctionTest, Renaming) {
const char* code = R"(
<
ir_version: 8,
opset_import: [ "" : 16, "local" : 1 ]
>
agraph (float[N] x) => (float[N] y)
{
y1 = local.myfun (x)
y = local.myfun (y1)
}
<
opset_import: [ "" : 16 ],
domain: "local"
>
myfun (lx) => (ly) {
two = Constant <value = float[1] {2.0}> ()
ly = Mul (lx, two)
}
)";
Check(code, "x", {1.0, 2.0, 3.0}, "y", {4.0, 8.0, 12.0});
}
} // namespace test
} // namespace onnxruntime