Copy input tensors (#395)

* Copy input tensors

* Check that default CPU execution provider is registered successfully

* Insert Memcpy only when an input is connected to both provider and non-provider nodes.
This commit is contained in:
Artem Rudoy 2019-02-01 14:53:45 -08:00 committed by GitHub
parent ebfed60741
commit 5cac965471
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 21 additions and 9 deletions

View file

@ -116,12 +116,23 @@ bool TransformerMemcpyImpl::ModifyGraph(const KernelRegistryManager& kernel_regi
// for initializers shared by different providers, create dups
ProcessInitializers();
for (auto arg : graph_.GetInputs())
BuildDefsMapping(arg, kernel_registries);
for (auto arg : non_provider_input_defs_)
BuildDefsMapping(arg, kernel_registries);
for (auto arg : non_provider_output_defs_)
BuildDefsMapping(arg, kernel_registries);
for (auto arg : graph_.GetInputs())
// For inputs we need to create a copy node only when the input is connected to both provider
// and non-provider nodes. Otherwise utils::CopyInputsAcrossDevices() will do the job.
if (provider_input_defs_.count(arg) && non_provider_input_defs_.count(arg)) {
AddCopyNode(const_cast<onnxruntime::NodeArg*>(arg), true);
modified = true;
}
for (auto arg : non_provider_output_defs_)
if (provider_input_defs_.count(arg)) {
AddCopyNode(arg, true);

View file

@ -327,8 +327,8 @@ class InferenceSession::Impl {
if (!execution_providers_.Get(onnxruntime::kCpuExecutionProvider)) {
LOGS(*session_logger_, INFO) << "Adding default CPU execution provider.";
CPUExecutionProviderInfo epi{session_options_.enable_cpu_mem_arena};
execution_providers_.Add(onnxruntime::kCpuExecutionProvider,
std::make_unique<CPUExecutionProvider>(epi));
ORT_RETURN_IF_ERROR(execution_providers_.Add(onnxruntime::kCpuExecutionProvider,
std::make_unique<CPUExecutionProvider>(epi)));
}
onnxruntime::Graph& graph = model_->MainGraph();

View file

@ -1000,17 +1000,17 @@ TEST(ExecutionProviderTest, FunctionTest) {
RunOptions run_options;
run_options.run_tag = so.session_logid;
CPUExecutionProviderInfo epi;
auto testCPUExecutionProvider = std::make_unique<::onnxruntime::CPUExecutionProvider>(epi);
std::vector<int64_t> dims_mul_x = {3, 2};
std::vector<float> values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
MLValue ml_value_x;
CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
&ml_value_x);
CreateMLValue<float>(testCPUExecutionProvider->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x, &ml_value_x);
MLValue ml_value_y;
CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
&ml_value_y);
CreateMLValue<float>(testCPUExecutionProvider->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x, &ml_value_y);
MLValue ml_value_z;
CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
&ml_value_z);
CreateMLValue<float>(testCPUExecutionProvider->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x, &ml_value_z);
NameMLValMap feeds;
feeds.insert(std::make_pair("X", ml_value_x));
feeds.insert(std::make_pair("Y", ml_value_y));
@ -1031,7 +1031,8 @@ TEST(ExecutionProviderTest, FunctionTest) {
VerifyOutputs(fetches, expected_dims_mul_m, expected_values_mul_m);
InferenceSession session_object_2{so};
session_object_2.RegisterExecutionProvider(std::make_unique<FuseExecutionProvider>());
session_object_2.RegisterExecutionProvider(std::move(testCPUExecutionProvider));
session_object_2.RegisterExecutionProvider(std::make_unique<::onnxruntime::FuseExecutionProvider>());
status = session_object_2.Load(model_file_name);
ASSERT_TRUE(status.IsOK());
status = session_object_2.Initialize();