feat(//py)!: Implementing top level python api changes to reflect new

Input type and enabled_precisions set

BREAKING CHANGE: This commit introduces the next iteration of the Python
TRTorch API. Starting in TRTorch v0.5.0 support for the "input_shapes"
and "op_precision" compile spec keys will be removed. Users should port
forward to using the "inputs" key which expects a list of trtorch.Input
objects and the "enabled_precisions" key which expects a set of data
type specifying enums.

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/setup.py

@@ -181,7 +181,7 @@ def run(self):
         include_dirs=[
             dir_path + "trtorch/csrc",
             dir_path + "/../",
-            dir_path + "/../bazel-TRTorch/external/tensorrt/include",
+            dir_path + "/../bazel-trtorch-testing/external/tensorrt/include",
         ],
         extra_compile_args=[
             "-Wno-deprecated",

py/trtorch/Input.py

@@ -0,0 +1,181 @@
+from enum import Enum
+from typing import List, Dict, Any
+
+import torch
+
+from trtorch import _types
+import trtorch._C
+
+class Input(object):
+    """
+    Defines an input to a module in terms of expected shape, data type and tensor format.
+
+    Attributes:
+        shape_mode (trtorch.Input._ShapeMode): Is input statically or dynamically shaped
+        shape (Tuple or Dict): Either a single Tuple or a dict of tuples defining the input shape.
+            Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form
+            ``{
+                "min_shape": Tuple,
+                "opt_shape": Tuple,
+                "max_shape": Tuple
+            }``
+        dtype (trtorch.dtype): The expected data type of the input tensor (default: trtorch.dtype.float32)
+        format (trtorch.TensorFormat): The expected format of the input tensor (default: trtorch.TensorFormat.NCHW)
+    """
+
+    class _ShapeMode(Enum):
+        STATIC = 0
+        DYNAMIC = 1
+
+    shape_mode = None
+    shape = None
+    dtype = _types.dtype.float32
+    format = _types.TensorFormat.contiguous
+
+    def __init__(self, *args, **kwargs):
+        """ __init__ Method for trtorch.Input
+
+        Input accepts one of a few construction patterns
+
+        Args:
+            shape (Tuple or List, optional): Static shape of input tensor
+
+        Keyword Arguments:
+            shape (Tuple or List, optional): Static shape of input tensor
+            min_shape (Tuple or List, optional): Min size of input tensor's shape range
+                Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implictly this sets Input's shape_mode to DYNAMIC
+            opt_shape (Tuple or List, optional): Opt size of input tensor's shape range
+                Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implictly this sets Input's shape_mode to DYNAMIC
+            max_shape (Tuple or List, optional): Max size of input tensor's shape range
+                Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implictly this sets Input's shape_mode to DYNAMIC
+            dtype (torch.dtype or trtorch.dtype): Expected data type for input tensor (default: trtorch.dtype.float32)
+            format (torch.memory_format or trtorch.TensorFormat): The expected format of the input tensor (default: trtorch.TensorFormat.NCHW)
+
+        Examples:
+            - Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)
+            - Input(shape=(1,3,32,32), dtype=trtorch.dtype.int32, format=trtorch.TensorFormat.NCHW)
+            - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=trtorch.dtype.float32, format=trtorch.TensorFormat.NCHW
+        """
+        if len(args) == 1:
+            if not Input._supported_input_size_type(args[0]):
+                raise TypeError(
+                "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                + str(type(args[0])))
+            if any(k in kwargs for k in ["min_shape", "opt_shape", "max_shape"]):
+                raise ValueError("Found that both shape (as a positional argument), and one or more of min_shape, opt_shape, max_shape were specified\nclass Input expects that only either shape or all three of min_shape, opt_shape, max_shape are defined")
+            self.shape = tuple(args[0])
+            self.shape_mode = Input._ShapeMode.STATIC
+
+        elif len(args) == 0:
+            if not ("shape" in kwargs) and not(all(k in kwargs for k in ["min_shape", "opt_shape", "max_shape"])):
+                raise ValueError("Missing required arguments for class Input\nEither shape or all three of min_shape, opt_shape, max_shape must be defined")
+            elif ("shape" in kwargs) and all(k in kwargs for k in ["min_shape", "opt_shape", "max_shape"]):
+                raise ValueError("Found that both shape, and one or more of min_shape, opt_shape, max_shape were specified\nclass Input expects that only either shape or all three of min_shape, opt_shape, max_shape are defined")
+
+            if "shape" in kwargs:
+                if not Input._supported_input_size_type(kwargs["shape"]):
+                    raise TypeError(
+                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                    + str(type(kwargs["shape"])))
+                self.shape = tuple(kwargs["shape"])
+                self.shape_mode = Input._ShapeMode.STATIC
+            else:
+                if not Input._supported_input_size_type(kwargs["min_shape"]):
+                    raise TypeError(
+                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                    + str(type(kwargs["min_shape"])) + " for min_shape")
+                if not Input._supported_input_size_type(kwargs["opt_shape"]):
+                    raise TypeError(
+                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                    + str(type(kwargs["opt_shape"])) + " for opt_shape")
+                if not Input._supported_input_size_type(kwargs["max_shape"]):
+                    raise TypeError(
+                    "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
+                    + str(type(kwargs["max_shape"])) + " for max_shape")
+
+                self.shape = {
+                    "min_shape": tuple(kwargs["min_shape"]),
+                    "opt_shape": tuple(kwargs["opt_shape"]),
+                    "max_shape": tuple(kwargs["max_shape"])
+                }
+                self.shape_mode = Input._ShapeMode.DYNAMIC
+
+            if "dtype" in kwargs:
+                self.dtype = Input._parse_dtype(kwargs["dtype"])
+
+            if "format" in kwargs:
+                self.format = Input._parse_format(kwargs["format"])
+
+        else:
+            raise ValueError("Unexpected number of positional arguments for class Input \n    Found {} arguments, expected either zero or a single positional arguments".format(len(args)))
+
+    def __str__(self) -> str:
+        if self.shape_mode == Input._ShapeMode.STATIC:
+            return "Input(shape={}, dtype={}, format={})".format(self.shape, str(self.dtype), str(self.format))
+        elif self.shape_mode == Input._ShapeMode.DYNAMIC:
+            return "Input(min_shape={}, opt_shape={}, max_shape={}, dtype={}, format={})".format(self.shape["min_shape"], self.shape["min_shape"], self.shape["min_shape"], str(self.dtype), str(self.format))
+        else:
+            raise RuntimeError("Unknown input shape mode")
+
+    def _to_internal(self) -> trtorch._C.Input:
+        internal_in = trtorch._C.Input()
+        if self.shape_mode == Input._ShapeMode.DYNAMIC:
+            internal_in.min = self.shape["min_shape"]
+            internal_in.opt = self.shape["opt_shape"]
+            internal_in.max = self.shape["max_shape"]
+            internal_in.input_is_dynamic = True
+        else:
+            internal_in.opt = self.shape
+            internal_in.input_is_dynamic = False
+        internal_in.dtype = self.dtype
+        internal_in.format = self.format
+        return internal_in
+
+    @staticmethod
+    def _supported_input_size_type(input_size: Any) -> bool:
+        if isinstance(input_size, torch.Size):
+            return True
+        elif isinstance(input_size, tuple):
+            return True
+        elif isinstance(input_size, list):
+            return True
+        else:
+            return False
+
+    @staticmethod
+    def _parse_dtype(dtype: Any) -> _types.dtype:
+        if isinstance(dtype, torch.dtype):
+            if dtype == torch.int32:
+                return _types.dtype.int32
+            elif dtype == torch.half:
+                return _types.dtype.half
+            elif dtype == torch.float:
+                return _types.dtype.float
+            elif dtype == torch.bool:
+                return _types.dtype.bool
+            else:
+                raise TypeError("Provided an unsupported data type as an input data type (support: bool, int32, half, float), got: " +
+                                str(dtype))
+
+        elif isinstance(dtype, _types.DataTypes):
+            return dtype
+
+        else:
+            raise TypeError("Input data type needs to be specified with a torch.dtype or a trtorch.dtype, got: " +
+                            str(type(dtype)))
+
+    @staticmethod
+    def _parse_format(format: Any) -> _types.TensorFormat:
+        if isinstance(format, torch.memory_format):
+            if format == torch.contiguous_format:
+                return _types.TensorFormat.contiguous
+            elif format == torch.channels_last:
+                return _types.TensorFormat.channel_last
+            else:
+                raise ValueError("Provided an unsupported tensor format (support: NHCW/contiguous_format, NHWC/channel_last)")
+
+        elif isinstance(format, _types.TensorFormat):
+            return format
+
+        else:
+            raise TypeError("Tensor format needs to be specified with either torch.memory_format or trtorch.TensorFormat")

py/trtorch/__init__.py

@@ -13,6 +13,7 @@
 from trtorch import ptq
 from trtorch._types import *
 from trtorch import logging
+from trtorch.Input import Input
 
 
 def _register_with_torch():

py/trtorch/_compile_spec.py

@@ -1,7 +1,10 @@
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Set
 import torch
 import trtorch._C
 from trtorch import _types
+from trtorch.Input import Input
+
+import warnings
 
 
 def _supported_input_size_type(input_size: Any) -> bool:
@@ -26,36 +29,23 @@ def _parse_input_ranges(input_sizes: List) -> List:
     for i in input_sizes:
         if isinstance(i, dict):
             if all(k in i for k in ["min", "opt", "min"]):
-                in_range = trtorch._C.InputRange()
-                in_range.min = i["min"]
-                in_range.opt = i["opt"]
-                in_range.max = i["max"]
-                parsed_input_sizes.append(in_range)
+                parsed_input_sizes.append(Input(min_shape=i["min"], opt_shape=i["opt"], max_shape=i["max"])._to_internal())
 
             elif "opt" in i:
-                in_range = trtorch._C.InputRange()
-                in_range.min = i["opt"]
-                in_range.opt = i["opt"]
-                in_range.max = i["opt"]
-                parsed_input_sizes.append(in_range)
+                parsed_input_sizes.append(Input(shape=i["opt"])._to_internal())
 
             else:
                 raise KeyError(
                     "An input size must either be a static size or a range of three sizes (min, opt, max) as Dict")
 
         elif isinstance(i, list):
-            in_range = trtorch._C.InputRange()
-            in_range.min = i
-            in_range.opt = i
-            in_range.max = i
-            parsed_input_sizes.append(in_range)
+            parsed_input_sizes.append(Input(shape=i)._to_internal())
 
         elif isinstance(i, tuple):
-            in_range = trtorch._C.InputRange()
-            in_range.min = list(i)
-            in_range.opt = list(i)
-            in_range.max = list(i)
-            parsed_input_sizes.append(in_range)
+            parsed_input_sizes.append(Input(shape=i)._to_internal())
+
+        elif isinstance(i, torch.Size):
+            parsed_input_sizes.append(Input(shape=i)._to_internal())
 
     return parsed_input_sizes
 
@@ -80,6 +70,15 @@ def _parse_op_precision(precision: Any) -> _types.dtype:
                         str(type(precision)))
 
 
+def _parse_enabled_precisions(precisions: Any) -> Set:
+    parsed_precisions = set()
+    if any([isinstance(precisions, type) for type in [list, tuple, set]]):
+        for p in precisions:
+            parsed_precisions.add(_parse_op_precision(p))
+    else:
+        parsed_precisions.add(_parse_op_precision(precisions))
+    return parsed_precisions
+
 def _parse_device_type(device: Any) -> _types.DeviceType:
     if isinstance(device, torch.device):
         if device.type == 'cuda':
@@ -140,39 +139,36 @@ def _parse_torch_fallback(fallback_info: Dict[str, Any]) -> trtorch._C.TorchFall
 
     return info
 
-def _parse_input_dtypes(input_dtypes: List) -> List:
-    parsed_input_dtypes = []
-    for dtype in input_dtypes:
-        if isinstance(dtype, torch.dtype):
-            if dtype == torch.int8:
-                parsed_input_dtypes.append(_types.dtype.int8)
-            elif dtype == torch.half:
-                parsed_input_dtypes.append(_types.dtype.half)
-            elif dtype == torch.float:
-                parsed_input_dtypes.append(_types.dtype.float)
-            elif dtype == torch.int32:
-                parsed_input_dtypes.append(_types.dtype.int32)
-            elif dtype == torch.bool:
-                parsed_input_dtypes.append(_types.dtype.bool)
-            else:
-                raise TypeError("Invalid input dtype. Supported input datatypes include float|half|int8|int32|bool), got: " + str(dtype))
-
-    return parsed_input_dtypes
-
 def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
     info = trtorch._C.CompileSpec()
-    if "input_shapes" not in compile_spec:
+    if "input_shapes" not in compile_spec and "inputs" not in compile_spec:
         raise KeyError(
-            "Input shapes for inputs are required as a List, provided as either a static sizes or a range of three sizes (min, opt, max) as Dict"
+            "Module input definitions are requried to compile module. Provide a list of trtorch.Input keyed to \"inputs\" in the compile spec"
         )
 
-    info.input_ranges = _parse_input_ranges(compile_spec["input_shapes"])
+    if "input_shapes" in compile_spec and "inputs" in compile_spec:
+        raise KeyError(
+            "Found both key \"input_shapes\", and \"inputs\" in compile spec, please port forward to using only \"inputs\""
+        )
+
+    if "input_shapes" in compile_spec:
+        warnings.warn("Key \"input_shapes\" is deprecated in favor of \"inputs\". Support for \"input_shapes\" will be removed in TRTorch v0.5.0", DeprecationWarning)
+        info.inputs = _parse_input_ranges(compile_spec["input_shapes"])
+
+    if "inputs" in compile_spec:
+        info.inputs = [ i._to_internal() for i in compile_spec["inputs"] ]
+
+    if "op_precision" in compile_spec and "enabled_precisions" in compile_spec:
+        raise KeyError(
+            "Found both key \"op_precision\", and \"enabled_precisions\" in compile spec, please port forward to using only \"enabled_precisions\""
+        )
 
     if "op_precision" in compile_spec:
-        info.op_precision = _parse_op_precision(compile_spec["op_precision"])
+        warnings.warn("Key \"op_precision\" is being deprecated in favor of \"enabled_precision\" which expects a set of precisions to be enabled during compilation (FP32 will always be enabled), Support for \"op_precision\" will be removed in TRTorch v0.5.0", DeprecationWarning)
+        info.enabled_precisions = _parse_enabled_precisions(compile_spec["op_precision"])
 
-    if "input_dtypes" in compile_spec:
-        info.input_dtypes = _parse_input_dtypes(compile_spec["input_dtypes"])
+    if "enabled_precisions" in compile_spec:
+        info.enabled_precisions = _parse_enabled_precisions(compile_spec["enabled_precision"])
 
     if "calibrator" in compile_spec:
         info.ptq_calibrator = compile_spec["calibrator"]
@@ -233,7 +229,8 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
     Args:
         compile_spec (dict): Compilation settings including operating precision, target device, etc.
             One key is required which is ``input_shapes``, describing the input sizes or ranges for inputs
-            to the graph. All other keys are optional. Entries for each method to be compiled.
+            to the graph as well as expect types and formats for those inputs. All other keys are optional.
+            Entries for each method to be compiled.
 
             Note: Partial compilation of TorchScript modules is not supported through the PyTorch TensorRT backend
             If you need this feature, use trtorch.compile to compile your module. Usage of the resulting module is
@@ -243,13 +240,15 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
 
                 CompileSpec = {
                     "forward" : trtorch.TensorRTCompileSpec({
-                        "input_shapes": [
-                            (1, 3, 224, 224), # Static input shape for input #1
-                            {
-                                "min": (1, 3, 224, 224),
-                                "opt": (1, 3, 512, 512),
-                                "max": (1, 3, 1024, 1024)
-                            } # Dynamic input shape for input #2
+                        "inputs": [
+                            trtorch.Input((1, 3, 224, 224)), # Static input shape for input #1
+                            trtorch.Input(
+                                min_shape=1, 3, 224, 224),
+                                opt_shape=(1, 3, 512, 512),
+                                max_shape=(1, 3, 1024, 1024),
+                                dtype=torch.int32
+                                format=torch.channel_last
+                            ) # Dynamic input shape for input #2
                         ],
                         "device": {
                             "device_type": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)
@@ -284,12 +283,15 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
 
     backend_spec = torch.classes.tensorrt.CompileSpec()
 
-    for i in parsed_spec.input_ranges:
-        ir = torch.classes.tensorrt._InputRange()
-        ir._set_min(i.min)
-        ir._set_opt(i.opt)
-        ir._set_max(i.max)
-        backend_spec._append_input_range(ir)
+    for i in parsed_spec.inputs:
+        clone = torch.classes.tensorrt._Input()
+        clone._set_min(i.min)
+        clone._set_opt(i.opt)
+        clone._set_max(i.max)
+        clone._set_dtype(i.dtype)
+        clone._set_format(i.format)
+        clone._set_input_is_dynamic(i.input_is_dynamic)
+        backend_spec._append_input(clone)
 
     d = torch.classes.tensorrt._Device()
     d._set_device_type(int(parsed_spec.device.device_type))
@@ -309,9 +311,7 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.
 
     backend_spec._set_device(d)
     backend_spec._set_torch_fallback(torch_fallback)
-    backend_spec._set_op_precision(int(parsed_spec.op_precision))
-    for dtype in parsed_spec.input_dtypes:
-        backend_spec._append_input_dtypes(int64_t(dtype))
+    backend_spec._set_precisions([int(i) for i in parsed_spec.enabled_precisions])
 
     backend_spec._set_disable_tf32(parsed_spec.disable_tf32)
     backend_spec._set_refit(parsed_spec.refit)