From 56b5a53059aecd9f8265e5677113986240d5dbae Mon Sep 17 00:00:00 2001 From: Bo Wang Date: Tue, 28 Jun 2022 13:15:20 -0700 Subject: [PATCH 1/2] fix: fix the parsing related model loading bug Signed-off-by: Bo Wang --- core/compiler.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/core/compiler.cpp b/core/compiler.cpp index 4a4389bea3..b154d52988 100644 --- a/core/compiler.cpp +++ b/core/compiler.cpp @@ -235,8 +235,11 @@ GraphAndMapping ConstructFallbackGraph( // the mapping from lowering graph => fallback global graph std::unordered_map old_to_new_g; - for (auto input : block->inputs()) { - util::getOrAddInputForValue(input, new_g, old_to_new_g); + + for (uint64_t i = 0; i < block->inputs().size(); i++) { + auto in_val = new_g->addInput(std::string("input_") + std::to_string(i)); + in_val->setType(block->inputs()[i]->type()); + old_to_new_g[block->inputs()[i]] = in_val; } for (auto& seg_block : segmented_blocks) { From e7c359d28af7f93cdf4974e6e3a75b7474229083 Mon Sep 17 00:00:00 2001 From: Bo Wang Date: Tue, 5 Jul 2022 15:35:17 -0700 Subject: [PATCH 2/2] test: add test case for parsing Signed-off-by: Bo Wang --- core/compiler.cpp | 13 ++++--- tests/core/partitioning/BUILD | 18 ++++++++- .../core/partitioning/test_loading_model.cpp | 39 +++++++++++++++++++ 3 files changed, 63 insertions(+), 7 deletions(-) create mode 100644 tests/core/partitioning/test_loading_model.cpp diff --git a/core/compiler.cpp b/core/compiler.cpp index b154d52988..c3493c9b84 100644 --- a/core/compiler.cpp +++ b/core/compiler.cpp @@ -235,11 +235,8 @@ GraphAndMapping ConstructFallbackGraph( // the mapping from lowering graph => fallback global graph std::unordered_map old_to_new_g; - - for (uint64_t i = 0; i < block->inputs().size(); i++) { - auto in_val = new_g->addInput(std::string("input_") + std::to_string(i)); - in_val->setType(block->inputs()[i]->type()); - old_to_new_g[block->inputs()[i]] = in_val; + for (auto input : block->inputs()) { + util::getOrAddInputForValue(input, new_g, old_to_new_g); } for (auto& seg_block : segmented_blocks) { @@ -439,7 +436,11 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg) auto graph_and_mapping = ConstructFallbackGraph(new_mod, g->block(), input_ivalues_map, cfg, static_params, fallback_nodes); new_g = graph_and_mapping.first; - LOG_INFO("Segmented Graph: " << *new_g); + // renaming the input name of graph after fallback to ensure pytorch deserialize it correctly + for (size_t i = 0; i < new_g->inputs().size(); ++i) { + new_g->inputs()[i]->setDebugName(std::string("input_") + std::to_string(i)); + } + LOG_INFO(*new_g << "(GraphAfterFallback)"); // if there is no tensorrt engine self in fallback graph, there is no conversion, we just return the initial // module diff --git a/tests/core/partitioning/BUILD b/tests/core/partitioning/BUILD index ec5e9c77fc..24d199a127 100644 --- a/tests/core/partitioning/BUILD +++ b/tests/core/partitioning/BUILD @@ -37,6 +37,21 @@ partitioning_test( name = "test_resolve_nontensor_inputs", ) +cc_test( + name = "test_loading_model", + srcs = ["test_loading_model.cpp"], + deps = [ + "//tests/util", + "@googletest//:gtest_main", + ] + select({ + ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + "//conditions:default": ["@libtorch//:libtorch"], + }), + data = [ + ":jit_models" + ] +) + cc_test( name = "test_fallback_graph_output", srcs = ["test_fallback_graph_output.cpp"], @@ -92,6 +107,7 @@ test_suite( ":test_fallback_graph_output", ":test_loop_fallback", ":test_conditionals", - ":test_resolve_nontensor_inputs" + ":test_resolve_nontensor_inputs", + ":test_loading_model" ] ) diff --git a/tests/core/partitioning/test_loading_model.cpp b/tests/core/partitioning/test_loading_model.cpp new file mode 100644 index 0000000000..057aaff2d8 --- /dev/null +++ b/tests/core/partitioning/test_loading_model.cpp @@ -0,0 +1,39 @@ +#include +#include +#include "core/compiler.h" +#include "gtest/gtest.h" +#include "tests/util/util.h" +#include "torch/script.h" + +#ifndef DISABLE_TEST_IN_CI + +TEST(Partitioning, ComputeResNet50FallbackGraphCorrectly) { + torch::jit::script::Module mod; + try { + mod = torch::jit::load("tests/modules/conditional_scripted.jit.pt"); + } catch (const c10::Error& e) { + std::cerr << "error loading the model\n"; + return; + } + + const std::vector> input_shapes = {{1, 3, 224, 224}}; + std::vector jit_inputs_ivalues; + std::vector trt_inputs_ivalues; + for (auto in_shape : input_shapes) { + auto in = at::randint(5, in_shape, {at::kCUDA}); + jit_inputs_ivalues.push_back(in.clone()); + trt_inputs_ivalues.push_back(in.clone()); + } + + std::vector input_ranges{torch_tensorrt::core::ir::Input({1, 3, 224, 224})}; + + torch_tensorrt::core::CompileSpec cfg(input_ranges); + cfg.partition_info.enabled = true; + + auto jit_results = mod.forward(jit_inputs_ivalues).toTensor(); + auto trt_mod = torch_tensorrt::core::CompileGraph(mod, cfg); + trt_mod.save("loading_model.ts"); + auto loaded_model = torch::jit::load("loading_model.ts"); +} + +#endif