feat: [collection] add unit test for complex collection model

Signed-off-by: inocsin <[email protected]>

Author: inocsin

tests/cpp/test_collection.cpp

@@ -353,4 +353,82 @@ TEST(CppAPITests, TestCollectionListInputOutput) {
 
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(out.toList().vec()[0].toTensor(), trt_out.toList().vec()[0].toTensor(), 1e-5));
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(out.toList().vec()[1].toTensor(), trt_out.toList().vec()[1].toTensor(), 1e-5));
+}
+
+
+TEST(CppAPITests, TestCollectionComplexModel) {
+
+  std::string path =
+  "/root/Torch-TensorRT/complex_model.ts";
+  torch::Tensor in0 = torch::randn({1, 3, 512, 512}, torch::kCUDA).to(torch::kHalf);
+  std::vector<at::Tensor> inputs;
+  inputs.push_back(in0);
+
+  torch::jit::Module mod;
+  try {
+    // Deserialize the ScriptModule from a file using torch::jit::load().
+    mod = torch::jit::load(path);
+  } catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n";
+  }
+  mod.eval();
+  mod.to(torch::kCUDA);
+  
+
+  std::vector<torch::jit::IValue> inputs_;
+
+  for (auto in : inputs) {
+    inputs_.push_back(torch::jit::IValue(in.clone()));
+  }
+
+  std::vector<torch::jit::IValue> complex_inputs;
+  auto input_list = c10::impl::GenericList(c10::TensorType::get());
+  input_list.push_back(inputs_[0]);
+  input_list.push_back(inputs_[0]);
+
+  torch::jit::IValue input_list_ivalue = torch::jit::IValue(input_list);
+
+  complex_inputs.push_back(input_list_ivalue);
+
+
+  auto out = mod.forward(complex_inputs);
+  LOG_DEBUG("Finish torchscirpt forward");
+
+
+  // auto input_shape = torch_tensorrt::Input(in0.sizes(), torch_tensorrt::DataType::kUnknown);
+  auto input_shape = torch_tensorrt::Input(in0.sizes(), torch_tensorrt::DataType::kHalf);
+
+  auto input_shape_ivalue = torch::jit::IValue(std::move(c10::make_intrusive<torch_tensorrt::Input>(input_shape)));
+
+
+  c10::TypePtr elementType = input_shape_ivalue.type();
+  auto list = c10::impl::GenericList(elementType);
+  list.push_back(input_shape_ivalue);
+  list.push_back(input_shape_ivalue);
+
+
+  torch::jit::IValue complex_input_shape(list);
+  std::tuple<torch::jit::IValue> input_tuple2(complex_input_shape);
+  torch::jit::IValue complex_input_shape2(input_tuple2);
+
+  auto compile_settings = torch_tensorrt::ts::CompileSpec(complex_input_shape2);
+  compile_settings.require_full_compilation = false;
+  compile_settings.min_block_size = 1;
+
+  // Need to skip the conversion of __getitem__ and ListConstruct
+  compile_settings.torch_executed_ops.push_back("aten::__getitem__");
+  compile_settings.torch_executed_ops.push_back("prim::ListConstruct");
+
+  // // FP16 execution
+  compile_settings.enabled_precisions = {torch::kHalf};
+  // // Compile module
+  auto trt_mod = torch_tensorrt::torchscript::compile(mod, compile_settings);
+  LOG_DEBUG("Finish compile");
+  auto trt_out = trt_mod.forward(complex_inputs);
+  // auto trt_out = trt_mod.forward(complex_inputs_list);
+
+  // std::cout << out.toTuple()->elements()[0].toTensor() << std::endl;
+  
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(out.toTuple()->elements()[0].toTensor(), trt_out.toTuple()->elements()[0].toTensor(), 1e-5));
+  ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(out.toTuple()->elements()[1].toTensor(), trt_out.toTuple()->elements()[1].toTensor(), 1e-5));
 }

tests/py/test_collection.py

@@ -79,6 +79,22 @@ def forward(self, z: List[torch.Tensor]):
         r = [r1, r2]
         return r
 
+class ComplexModel(nn.Module):
+    def __init__(self):
+        super(ComplexModel, self).__init__()
+        self.list_model = ListInputOutput()
+        self.tuple_model = TupleInputOutput()
+
+    def forward(self, z: List[torch.Tensor]):
+        r1 = z[0] + z[1]
+        r2 = z[0] - z[1]
+        r3 = (r1, r2)
+        r4 = [r2, r1]
+        tuple_out = self.tuple_model(r3)
+        list_out = self.list_model(r4)
+        r = (tuple_out[1], list_out[0])
+        return r
+
 input_data = torch.randn((16, 3, 32, 32))
 input_data = input_data.float().to("cuda")
 
@@ -115,4 +131,11 @@ def forward(self, z: List[torch.Tensor]):
 print(list_input_ts.graph)
 result = list_input_ts([input_data, input_data])
 list_input_ts.to("cuda").eval()
-torch.jit.save(list_input_ts, "./list_input_output.ts")
+torch.jit.save(list_input_ts, "./list_input_output.ts")
+
+complex_model = ComplexModel()
+complex_model_ts = torch.jit.script(complex_model)
+print(complex_model_ts.graph)
+result = complex_model_ts([input_data, input_data])
+complex_model_ts.to("cuda").eval()
+torch.jit.save(complex_model_ts, "./complex_model.ts")