Repository: pytorch/TensorRT

Issue:
API to register arbitrary TRT engines as modules
It would be nice to be able to accept TRT engines built outside TRTorch and register them in the execution system. 


Discussion:

Resolving patch:
diff --git a/py/trtorch/_compiler.py b/py/trtorch/_compiler.py
--- a/py/trtorch/_compiler.py
+++ b/py/trtorch/_compiler.py
@@ -124,6 +124,26 @@ def convert_method_to_trt_engine(module: torch.jit.ScriptModule, method_name: st
     return trtorch._C.convert_graph_to_trt_engine(module._c, method_name, _parse_compile_spec(compile_spec))
 
 
+def embed_engine_in_new_module(serialized_engine: bytes) -> torch.jit.ScriptModule:
+    """Takes a pre-built serialized TensorRT engine and embeds it within a TorchScript module
+
+    Takes a pre-built serialied TensorRT engine (as bytes) and embeds it within a TorchScript module.
+    Registers the forward method to execute the TensorRT engine with the function signature:
+
+        forward(Tensor[]) -> Tensor[]
+
+    Module can be save with engine embedded with torch.jit.save and moved / loaded according to TRTorch portability rules
+
+    Args:
+        serialized_engine (bytes): Serialized TensorRT engine from either TRTorch or TensorRT APIs
+
+    Returns:
+        torch.jit.ScriptModule: New TorchScript module with engine embedded
+    """
+    cpp_mod = trtorch._C.embed_engine_in_new_module(serialized_engine)
+    return torch.jit._recursive.wrap_cpp_module(cpp_mod)
+
+
 def check_method_op_support(module: torch.jit.ScriptModule, method_name: str) -> bool:
     """Checks to see if a method is fully supported by TRTorch
 


Test patch:
diff --git a/tests/modules/test_modules_as_engines.cpp b/tests/modules/test_modules_as_engines.cpp
--- a/tests/modules/test_modules_as_engines.cpp
+++ b/tests/modules/test_modules_as_engines.cpp
@@ -16,6 +16,34 @@ TEST_P(ModuleTests, ModuleAsEngineIsClose) {
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-5));
 }
 
+TEST_P(ModuleTests, ModuleToEngineToModuleIsClose) {
+  std::vector<at::Tensor> inputs;
+  std::vector<torch::jit::IValue> inputs_ivalues;
+  for (auto in_shape : input_shapes) {
+    inputs.push_back(at::randint(5, in_shape, {at::kCUDA}));
+    inputs_ivalues.push_back(inputs[inputs.size() - 1].clone());
+  }
+
+  torch::jit::IValue jit_results_ivalues = trtorch::tests::util::RunModuleForward(mod, inputs_ivalues);
+  std::vector<at::Tensor> jit_results;
+  jit_results.push_back(jit_results_ivalues.toTensor());
+
+  auto forward_graph = mod.get_method("forward");
+  std::vector<c10::ArrayRef<int64_t>> input_ranges;
+  for (auto in : inputs) {
+    input_ranges.push_back(in.sizes());
+  }
+
+  auto engine = trtorch::ConvertGraphToTRTEngine(mod, "forward", input_ranges);
+  auto trt_mod = trtorch::EmbedEngineInNewModule(engine);
+
+  torch::jit::IValue trt_results_ivalues = trtorch::tests::util::RunModuleForward(mod, inputs_ivalues);
+  std::vector<at::Tensor> trt_results;
+  trt_results.push_back(trt_results_ivalues.toTensor());
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-5));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     ModuleAsEngineForwardIsCloseSuite,
     ModuleTests,
diff --git a/tests/py/BUILD b/tests/py/BUILD
--- a/tests/py/BUILD
+++ b/tests/py/BUILD
@@ -30,7 +30,7 @@ py_test(
     srcs = [
         "test_ptq_dataloader_calibrator.py",
         "model_test_case.py"
-    ]
+    ],
     deps = [
         requirement("torchvision")
     ]
@@ -43,7 +43,7 @@ py_test(
     srcs = [
         "test_ptq_trt_calibrator.py",
         "model_test_case.py"
-    ]
+    ],
     deps = [
         requirement("torchvision")
     ]
@@ -56,8 +56,6 @@ py_test(
         "test_multi_gpu.py",
         "model_test_case.py"
     ],
-    "//conditions:default" : []
-    }),
     deps = [
         requirement("torchvision")
     ]
@@ -74,12 +72,23 @@ py_test(
     ]
 )
 
+py_test(
+    name = "test_trt_intercompatability",
+    srcs = [
+        "test_trt_intercompatability.py",
+        "model_test_case.py"
+    ],
+    deps = [
+        requirement("torchvision")
+    ]
+)
+
 py_test(
     name = "test_ptq_to_backend",
     srcs = [
         "test_ptq_to_backend.py",
         "model_test_case.py"
-    ]
+    ],
     deps = [
         requirement("torchvision")
     ]
diff --git a/tests/py/test_api.py b/tests/py/test_api.py
--- a/tests/py/test_api.py
+++ b/tests/py/test_api.py
@@ -46,6 +46,30 @@ def test_compile_script(self):
         self.assertTrue(same < 2e-3)
 
 
+class TestPTtoTRTtoPT(ModelTestCase):
+
+    def setUp(self):
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+        self.ts_model = torch.jit.script(self.model)
+
+    def test_pt_to_trt_to_pt(self):
+        compile_spec = {
+            "input_shapes": [self.input.shape],
+            "device": {
+                "device_type": trtorch.DeviceType.GPU,
+                "gpu_id": 0,
+                "dla_core": 0,
+                "allow_gpu_fallback": False,
+                "disable_tf32": False
+            }
+        }
+
+        trt_engine = trtorch.convert_method_to_trt_engine(self.ts_model, "forward", compile_spec)
+        trt_mod = trtorch.embed_engine_in_new_module(trt_engine)
+        same = (trt_mod(self.input) - self.ts_model(self.input)).abs().max()
+        self.assertTrue(same < 2e-3)
+
+
 class TestCheckMethodOpSupport(unittest.TestCase):
 
     def setUp(self):
@@ -59,13 +83,13 @@ def test_check_support(self):
 class TestLoggingAPIs(unittest.TestCase):
 
     def test_logging_prefix(self):
-        new_prefix = "TEST"
+        new_prefix = "Python API Test: "
         trtorch.logging.set_logging_prefix(new_prefix)
         logging_prefix = trtorch.logging.get_logging_prefix()
         self.assertEqual(new_prefix, logging_prefix)
 
     def test_reportable_log_level(self):
-        new_level = trtorch.logging.Level.Warning
+        new_level = trtorch.logging.Level.Error
         trtorch.logging.set_reportable_log_level(new_level)
         level = trtorch.logging.get_reportable_log_level()
         self.assertEqual(new_level, level)
@@ -78,10 +102,11 @@ def test_is_colored_output_on(self):
 
 def test_suite():
     suite = unittest.TestSuite()
+    suite.addTest(unittest.makeSuite(TestLoggingAPIs))
     suite.addTest(TestCompile.parametrize(TestCompile, model=models.resnet18(pretrained=True)))
     suite.addTest(TestCompile.parametrize(TestCompile, model=models.mobilenet_v2(pretrained=True)))
+    suite.addTest(TestPTtoTRTtoPT.parametrize(TestPTtoTRTtoPT, model=models.mobilenet_v2(pretrained=True)))
     suite.addTest(unittest.makeSuite(TestCheckMethodOpSupport))
-    suite.addTest(unittest.makeSuite(TestLoggingAPIs))
 
     return suite
 
diff --git a/tests/py/test_trt_intercompatability.py b/tests/py/test_trt_intercompatability.py
new file mode 100644
--- /dev/null
+++ b/tests/py/test_trt_intercompatability.py
@@ -0,0 +1,55 @@
+import unittest
+import trtorch
+import torch
+import torchvision.models as models
+import tensorrt as trt
+
+from model_test_case import ModelTestCase
+
+
+class TestPyTorchToTRTEngine(ModelTestCase):
+
+    def setUp(self):
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda:0")
+        self.ts_model = torch.jit.script(self.model)
+
+    def test_pt_to_trt(self):
+        compile_spec = {
+            "input_shapes": [self.input.shape],
+            "device": {
+                "device_type": trtorch.DeviceType.GPU,
+                "gpu_id": 0,
+                "dla_core": 0,
+                "allow_gpu_fallback": False,
+                "disable_tf32": False
+            }
+        }
+
+        trt_engine = trtorch.convert_method_to_trt_engine(self.ts_model, "forward", compile_spec)
+
+        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+        with trt.Runtime(TRT_LOGGER) as rt:
+            engine = rt.deserialize_cuda_engine(trt_engine)
+            with engine.create_execution_context() as ctx:
+                out = torch.empty(size=tuple(engine.get_binding_shape(1))).to("cuda:0")
+                bindings = [self.input.contiguous().data_ptr(), out.contiguous().data_ptr()]
+                ctx.execute_async(batch_size=1,
+                                  bindings=bindings,
+                                  stream_handle=torch.cuda.current_stream(device='cuda:0').cuda_stream)
+                same = (out - self.ts_model(self.input)).abs().max()
+                self.assertTrue(same < 2e-3)
+
+
+def test_suite():
+    suite = unittest.TestSuite()
+    suite.addTest(TestPyTorchToTRTEngine.parametrize(TestPyTorchToTRTEngine, model=models.resnet18(pretrained=True)))
+
+    return suite
+
+
+suite = test_suite()
+
+runner = unittest.TextTestRunner()
+result = runner.run(suite)
+
+exit(int(not result.wasSuccessful()))