fix: Address review comments

gs-olive · gs-olive · commit 366cd311cce1 · 2023-07-17T23:10:15.000-07:00
- Ensure imports of utilities reference new directory structure
- Update the test cases to reflect the changes to `prepare_inputs.py`
- Add non-breaking functionality to `_Input` class
- Rename Python TRT runtime
- Reword runtime detection log message
diff --git a/py/torch_tensorrt/_Input.py b/py/torch_tensorrt/_Input.py
@@ -302,47 +302,58 @@ def _parse_tensor_domain(domain: Optional[Tuple[float, float]]) -> Tuple:
         return result_domain
 
     @classmethod
-    def from_tensor(cls, t: torch.Tensor) -> "Input":
+    def from_tensor(
+        cls, t: torch.Tensor, disable_memory_format_check: bool = False
+    ) -> "Input":
         """
         Produce a Input which contains the information of the given PyTorch tensor.
 
         Args:
             tensor (torch.Tensor): A PyTorch tensor.
+            disable_memory_format_check (bool): Whether to validate the memory formats of input tensors
 
         Returns:
             A Input object.
         """
-        if not any(
-            [
-                t.is_contiguous(memory_format=torch.contiguous_format),
-                t.is_contiguous(memory_format=torch.channels_last),
-            ]
+        if not (
+            t.is_contiguous(memory_format=torch.contiguous_format)
+            or t.is_contiguous(memory_format=torch.channels_last)
+            or disable_memory_format_check
         ):
             raise ValueError(
                 "Tensor does not have a supported memory format, supported formats are contiguous or channel_last"
             )
         frmt = (
             torch.contiguous_format
-            if t.is_contiguous(memory_format=torch.contiguous_format)
+            if (
+                t.is_contiguous(memory_format=torch.contiguous_format)
+                or disable_memory_format_check
+            )
             else torch.channels_last
         )
         return cls(shape=t.shape, dtype=t.dtype, format=frmt)
 
     @classmethod
-    def from_tensors(cls, ts: torch.Tensor) -> List["Input"]:
+    def from_tensors(
+        cls, ts: torch.Tensor, disable_memory_format_check: bool = False
+    ) -> List["Input"]:
         """
         Produce a list of Inputs which contain
         the information of all the given PyTorch tensors.
 
         Args:
             tensors (Iterable[torch.Tensor]): A list of PyTorch tensors.
+            disable_memory_format_check (bool): Whether to validate the memory formats of input tensors
 
         Returns:
             A list of Inputs.
         """
 
         assert isinstance(ts, (list, tuple))
-        return [cls.from_tensor(t) for t in ts]
+        return [
+            cls.from_tensor(t, disable_memory_format_check=disable_memory_format_check)
+            for t in ts
+        ]
 
     def example_tensor(self, optimization_profile_field: str = None) -> torch.Tensor:
         """
diff --git a/py/torch_tensorrt/dynamo/conversion/conversion.py b/py/torch_tensorrt/dynamo/conversion/conversion.py
@@ -1,7 +1,7 @@
 from typing import Sequence, Union
 import torch
 import io
-from torch_tensorrt.dynamo.runtime import TRTModule
+from torch_tensorrt.dynamo.runtime import _PythonTorchTRTModule
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt import Input
 from torch_tensorrt.dynamo.conversion import TRTInterpreter
@@ -23,7 +23,7 @@ def convert_module(
         settings: Compilation settings
         name: TRT engine name
     Returns:
-        TRTModule or TRTModuleNext
+        _PythonTorchTRTModule or TorchTensorRTModule
     """
     # Specify module output data types to ensure TRT output types agree with
     # that of the equivalent Torch module
@@ -35,7 +35,7 @@ def convert_module(
     output_dtypes = list(output.dtype for output in module_outputs)
     interpreter = TRTInterpreter(
         module,
-        Input.from_tensors(inputs),
+        Input.from_tensors(inputs, disable_memory_format_check=True),
         logger_level=(trt.Logger.VERBOSE if settings.debug else trt.Logger.WARNING),
         output_dtypes=output_dtypes,
     )
@@ -53,7 +53,7 @@ def convert_module(
     )
 
     if settings.use_python_runtime:
-        return TRTModule(
+        return _PythonTorchTRTModule(
             engine=interpreter_result.engine,
             input_names=interpreter_result.input_names,
             output_names=interpreter_result.output_names,
diff --git a/py/torch_tensorrt/dynamo/runtime/_PythonTorchTRTModule.py b/py/torch_tensorrt/dynamo/runtime/_PythonTorchTRTModule.py
@@ -6,8 +6,8 @@
 from torch_tensorrt.fx.utils import unified_dtype_converter, Frameworks
 
 
-class TRTModule(torch.nn.Module):
-    """TRTModule is a PyTorch module which encompasses an arbitrary TensorRT Engine.
+class PythonTorchTRTModule(torch.nn.Module):
+    """PythonTorchTRTModule is a PyTorch module which encompasses an arbitrary TensorRT Engine.
 
     This module is backed by the Torch-TensorRT runtime and is only compatibile with
     FX / Dynamo / Python deployments. This module cannot be serialized to torchscript via torch.jit.trace for C++ deployment.
@@ -16,8 +16,8 @@ class TRTModule(torch.nn.Module):
     def __init__(
         self, engine=None, input_names=None, output_names=None, cuda_graph_batch_size=-1
     ):
-        super(TRTModule, self).__init__()
-        self._register_state_dict_hook(TRTModule._on_state_dict)
+        super(PythonTorchTRTModule, self).__init__()
+        self._register_state_dict_hook(PythonTorchTRTModule._on_state_dict)
         self.engine = engine
         self.input_names = input_names
         self.output_names = output_names
@@ -94,7 +94,7 @@ def _initialize(self):
 
     def _check_initialized(self):
         if not self.initialized:
-            raise RuntimeError("TRTModule is not initialized.")
+            raise RuntimeError("PythonTorchTRTModule is not initialized.")
 
     def _on_state_dict(self, state_dict, prefix, local_metadata):
         self._check_initialized()
@@ -138,10 +138,12 @@ def __setstate__(self, state):
             self.context = self.engine.create_execution_context()
 
     def forward(self, *inputs):
-        with torch.autograd.profiler.record_function("TRTModule:Forward"):
+        with torch.autograd.profiler.record_function("PythonTorchTRTModule:Forward"):
             self._check_initialized()
 
-            with torch.autograd.profiler.record_function("TRTModule:ProcessInputs"):
+            with torch.autograd.profiler.record_function(
+                "PythonTorchTRTModule:ProcessInputs"
+            ):
                 assert len(inputs) == len(
                     self.input_names
                 ), f"Wrong number of inputs, expect {len(self.input_names)} get {len(inputs)}."
@@ -176,7 +178,9 @@ def forward(self, *inputs):
                             f"Expect {self.input_shapes[i]}, got {inputs[i].size()[1:]}."
                         )
 
-            with torch.autograd.profiler.record_function("TRTModule:ProcessOutputs"):
+            with torch.autograd.profiler.record_function(
+                "PythonTorchTRTModule:ProcessOutputs"
+            ):
                 # create output tensors
                 outputs: List[torch.Tensor] = []
 
@@ -207,7 +211,9 @@ def forward(self, *inputs):
                     )
                     bindings[idx] = output.data_ptr()
 
-            with torch.autograd.profiler.record_function("TRTModule:TensorRTRuntime"):
+            with torch.autograd.profiler.record_function(
+                "PythonTorchTRTModule:TensorRTRuntime"
+            ):
                 if self.engine.has_implicit_batch_dimension:
                     self.context.execute_async(
                         batch_size, bindings, torch.cuda.current_stream().cuda_stream
diff --git a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py
@@ -36,7 +36,7 @@ def __init__(
         output_binding_names: List[str] = [],
         target_device: Device = Device._current_device(),
     ):
-        """__init__ method for torch_tensorrt.TorchTensorRTModule
+        """__init__ method for torch_tensorrt.dynamo.runtime._TorchTensorRTModule.TorchTensorRTModule
 
         Takes a name, target device, serialized TensorRT engine, and binding names / order and constructs
         a PyTorch ``torch.nn.Module`` around it.
@@ -61,11 +61,11 @@ def __init__(
                     engine_bytes.write(trt_engine.serialize())
                     engine_str = engine_bytes.getvalue()
 
-                trt_module = TRTModule(
+                trt_module = TorchTensorRTModule(
                     engine_str,
-                    engine_name="my_module",
-                    input_names=["x"],
-                    output_names=["output"],
+                    name="my_module",
+                    input_binding_names=["x"],
+                    output_binding_names=["output"],
                 )
 
         """
diff --git a/py/torch_tensorrt/dynamo/runtime/__init__.py b/py/torch_tensorrt/dynamo/runtime/__init__.py
@@ -1,2 +1,2 @@
-from ._PythonTorchTRTModule import TRTModule
+from ._PythonTorchTRTModule import PythonTorchTRTModule
 from ._TorchTensorRTModule import TorchTensorRTModule
diff --git a/py/torch_tensorrt/dynamo/utils.py b/py/torch_tensorrt/dynamo/utils.py
@@ -36,7 +36,7 @@ def use_python_runtime_parser(use_python_runtime: Optional[bool] = None) -> bool
             reason = "since import failed, C++ dependency not installed"
 
     logger.info(
-        f"Using {'Python' if using_python_runtime else 'C++'} {reason} TRT Runtime"
+        f"Using {'Python-only' if using_python_runtime else 'Default'} Torch-TRT Runtime ({reason})"
     )
 
     return using_python_runtime
diff --git a/tests/py/dynamo/backend/test_compiler_utils.py b/tests/py/dynamo/backend/test_compiler_utils.py
@@ -24,16 +24,22 @@ def test_prepare_trt_device(self):
 class TestPrepareInputs(unittest.TestCase):
     def test_prepare_single_tensor_input(self):
         inputs = [torch.ones((4, 4))]
-        prepared_inputs = prepare_inputs(inputs)
+        prepared_inputs_trt, prepared_inputs_torch = prepare_inputs(inputs)
         self.assertTrue(
-            same_output_format(inputs, prepared_inputs, enforce_tensor_type=False)
+            same_output_format(inputs, prepared_inputs_trt, enforce_tensor_type=False)
+        )
+        self.assertTrue(
+            same_output_format(inputs, prepared_inputs_torch, enforce_tensor_type=False)
         )
 
     def test_prepare_trt_input(self):
         inputs = [torch_tensorrt.Input(shape=(4, 3), dtype=torch.float)]
-        prepared_inputs = prepare_inputs(inputs)
+        prepared_inputs_trt, prepared_inputs_torch = prepare_inputs(inputs)
+        self.assertTrue(
+            same_output_format(inputs, prepared_inputs_trt, enforce_tensor_type=False)
+        )
         self.assertTrue(
-            same_output_format(inputs, prepared_inputs, enforce_tensor_type=False)
+            same_output_format(inputs, prepared_inputs_torch, enforce_tensor_type=False)
         )
 
     def test_prepare_mixed_type_compound_tensor_input(self):
@@ -47,9 +53,12 @@ def test_prepare_mixed_type_compound_tensor_input(self):
                 (torch.rand((5, 1)), torch_tensorrt.Input(shape=(2, 3))),
             ),
         }
-        prepared_inputs = prepare_inputs(inputs)
+        prepared_inputs_trt, prepared_inputs_torch = prepare_inputs(inputs)
+        self.assertTrue(
+            same_output_format(inputs, prepared_inputs_trt, enforce_tensor_type=False)
+        )
         self.assertTrue(
-            same_output_format(inputs, prepared_inputs, enforce_tensor_type=False)
+            same_output_format(inputs, prepared_inputs_torch, enforce_tensor_type=False)
         )
 
 
diff --git a/tests/py/ts/api/test_classes.py b/tests/py/ts/api/test_classes.py
@@ -1,6 +1,6 @@
 import unittest
 import torch_tensorrt as torchtrt
-from torch_tensorrt.dynamo._TorchTensorRTModule import TorchTensorRTModule
+from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule
 import torch
 import torchvision.models as models
 import copy

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-from ._PythonTorchTRTModule import TRTModule`
	`1`	`+from ._PythonTorchTRTModule import PythonTorchTRTModule`
`2`	`2`	`from ._TorchTensorRTModule import TorchTensorRTModule`
Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ def use_python_runtime_parser(use_python_runtime: Optional[bool] = None) -> bool`
`36`	`36`	`reason = "since import failed, C++ dependency not installed"`
`37`	`37`
`38`	`38`	`logger.info(`
`39`		`- f"Using {'Python' if using_python_runtime else 'C++'} {reason} TRT Runtime"`
	`39`	`+ f"Using {'Python-only' if using_python_runtime else 'Default'} Torch-TRT Runtime ({reason})"`
`40`	`40`	`)`
`41`	`41`
`42`	`42`	`return using_python_runtime`