Merge pull request #15 from bddppq/misc-fixes

Fix indeterministic params order in `RunGraph`

Author: narendasan

tests/core/converters/test_linear.cpp

@@ -8,7 +8,7 @@
 TEST(Converters, ATenLinearNoBiasConvertsCorrectly) {
     const auto graph = R"IR(
       graph(%0 : Tensor,
-            %1 : Float(2, 2)):
+            %1 : Float(3, 2)):
         %2 : None = prim::Constant()
         %3 : Tensor = aten::linear(%0, %1, %2)
         return (%3))IR";
@@ -19,7 +19,7 @@ TEST(Converters, ATenLinearNoBiasConvertsCorrectly) {
     //Input Tensor needs to be 4D for TensorRT linear
     auto in = at::randint(1, 10, {1, 2}, {at::kCUDA});
     auto w = at::randint(1, 10, {3, 2}, {at::kCUDA});
-    
+
     auto params = trtorch::core::conversion::get_named_params(g->inputs(), {w});
     auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
 
@@ -32,8 +32,6 @@ TEST(Converters, ATenLinearNoBiasConvertsCorrectly) {
 }
 
 
-//TODO: Track down the cause of why the JIT linear function fails
-//TODO: Sort out what the exepected output dim should be? 
 TEST(Converters, ATenLinearBiasConvertsCorrectly) {
     const auto graph = R"IR(
       graph(%0 : Tensor,
@@ -53,7 +51,7 @@ TEST(Converters, ATenLinearBiasConvertsCorrectly) {
     auto jit_in = at::clone(in);
     auto jit_w = at::clone(w);
     auto jit_b = at::clone(b);
-    
+
     auto params = trtorch::core::conversion::get_named_params(g->inputs(), {jit_w, jit_b});
     auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
 

tests/models/BUILD

@@ -15,10 +15,7 @@ cc_test(
     name = "test_multiple_registered_engines",
     srcs = ["test_multiple_registered_engines.cpp"],
     deps = [
-        "//cpp/api:trtorch",
-        "//tests/util",
-        "@libtorch//:libtorch",
-        "@googletest//:gtest_main",
+        ":module_test",
     ],
     data = [
         ":jit_models"

tests/models/test_multiple_registered_engines.cpp

@@ -1,8 +1,8 @@
 #include <string>
 #include "gtest/gtest.h"
-#include "torch/csrc/jit/irparser.h"
+#include "torch/script.h"
 #include "tests/util/util.h"
-#include "cpp/trtorch.h"
+#include "trtorch/trtorch.h"
 
 TEST(ModuleTests, CanRunMultipleEngines) {
     torch::jit::script::Module mod1;
@@ -16,7 +16,7 @@ TEST(ModuleTests, CanRunMultipleEngines) {
         return;
     }
 
-    const std::vector<int64_t> input_shape = {1,3,224,224};
+    const std::vector<std::vector<int64_t>> input_shapes = {{1,3,224,224}};
 
     std::vector<torch::jit::IValue> jit1_inputs_ivalues;
     std::vector<torch::jit::IValue> trt1_inputs_ivalues;
@@ -38,18 +38,18 @@ TEST(ModuleTests, CanRunMultipleEngines) {
     std::vector<at::Tensor> jit1_results;
     jit1_results.push_back(jit1_results_ivalues.toTensor());
 
-        torch::jit::IValue jit2_results_ivalues = trtorch::tests::util::RunModuleForward(mod2, jit2_inputs_ivalues);
+    torch::jit::IValue jit2_results_ivalues = trtorch::tests::util::RunModuleForward(mod2, jit2_inputs_ivalues);
     std::vector<at::Tensor> jit2_results;
     jit2_results.push_back(jit2_results_ivalues.toTensor());
 
 
     auto trt_mod1 = trtorch::CompileGraph(mod1, input_shapes);
-    torch::jit::IValue trt1_results_ivalues = trtorch::tests::util::RunModuleForward(trt1_mod, trt1_inputs_ivalues);
+    torch::jit::IValue trt1_results_ivalues = trtorch::tests::util::RunModuleForward(trt_mod1, trt1_inputs_ivalues);
     std::vector<at::Tensor> trt1_results;
     trt1_results.push_back(trt1_results_ivalues.toTensor());
 
      auto trt_mod2 = trtorch::CompileGraph(mod2, input_shapes);
-    torch::jit::IValue trt2_results_ivalues = trtorch::tests::util::RunModuleForward(trt2_mod, trt2_inputs_ivalues);
+    torch::jit::IValue trt2_results_ivalues = trtorch::tests::util::RunModuleForward(trt_mod2, trt2_inputs_ivalues);
     std::vector<at::Tensor> trt2_results;
     trt2_results.push_back(trt2_results_ivalues.toTensor());
 

tests/util/run_graph.cpp

@@ -21,8 +21,11 @@ std::vector<at::Tensor> RunGraph(std::shared_ptr<torch::jit::Graph>& g,
         inputs_.push_back(in.clone());
     }
 
-    for (auto p : params) {
-        inputs_.push_back(p.second.clone());
+    for (auto* in : g->inputs()) {
+        const auto iter = params.find(in);
+        if (iter != params.end()) {
+            inputs_.push_back(iter->second.clone());
+        }
     }
 
     torch::jit::GraphExecutor executor(g);