test: Add class test for torch_tensorrt.Input

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

Author: narendasan

py/torch_tensorrt/_Input.py

@@ -214,7 +214,7 @@ def _parse_format(format: Any) -> _enums.TensorFormat:
             if format == torch.contiguous_format:
                 return _enums.TensorFormat.contiguous
             elif format == torch.channels_last:
-                return _enums.TensorFormat.channel_last
+                return _enums.TensorFormat.channels_last
             else:
                 raise ValueError(
                     "Provided an unsupported tensor format (support: NHCW/contiguous_format, NHWC/channel_last)")

py/torch_tensorrt/csrc/tensorrt_classes.cpp

@@ -49,7 +49,7 @@ Device::Device(const core::runtime::CudaDevice& internal_dev) {
 
 nvinfer1::TensorFormat toTRTTensorFormat(TensorFormat value) {
   switch (value) {
-    case TensorFormat::kChannelLast:
+    case TensorFormat::kChannelsLast:
       return nvinfer1::TensorFormat::kHWC;
     case TensorFormat::kContiguous:
     default:
@@ -61,7 +61,7 @@ std::string to_str(TensorFormat value) {
   switch (value) {
     case TensorFormat::kContiguous:
       return "Contiguous/Linear/NCHW";
-    case TensorFormat::kChannelLast:
+    case TensorFormat::kChannelsLast:
       return "Channel Last/NHWC";
     default:
       return "UNKNOWN";

py/torch_tensorrt/csrc/tensorrt_classes.h

@@ -31,7 +31,7 @@ enum class DataType : int8_t { kFloat, kHalf, kChar, kInt32, kBool, kUnknown };
 std::string to_str(DataType value);
 nvinfer1::DataType toTRTDataType(DataType value);
 
-enum class TensorFormat : int8_t { kContiguous, kChannelLast };
+enum class TensorFormat : int8_t { kContiguous, kChannelsLast };
 std::string to_str(TensorFormat value);
 nvinfer1::TensorFormat toTRTTensorFormat(TensorFormat value);
 

py/torch_tensorrt/csrc/torch_tensorrt_py.cpp

@@ -204,7 +204,7 @@ PYBIND11_MODULE(_C, m) {
 
   py::enum_<TensorFormat>(m, "TensorFormat", "Enum to specifiy the memory layout of tensors")
       .value("contiguous", TensorFormat::kContiguous, "Contiguous memory layout (NCHW / Linear)")
-      .value("channel_last", TensorFormat::kChannelLast, "Channel last memory layout (NHWC)")
+      .value("channels_last", TensorFormat::kChannelsLast, "Channels last memory layout (NHWC)")
       .export_values();
 
   py::enum_<nvinfer1::CalibrationAlgoType>(m, "CalibrationAlgo", py::module_local(), "Type of calibration algorithm")

tests/py/test_api.py

@@ -3,6 +3,7 @@
 import torch
 import torchvision.models as models
 import copy
+from typing import Dict
 
 from model_test_case import ModelTestCase
 
@@ -351,6 +352,138 @@ def test_from_torch(self):
         self.assertEqual(device.device_type, torchtrt.DeviceType.GPU)
         self.assertEqual(device.gpu_id, 0)
 
+class TestInput(unittest.TestCase):
+
+    def _verify_correctness(self, struct: torchtrt.Input, target: Dict) -> bool:
+        internal = struct._to_internal()
+
+        list_eq = lambda al, bl: all([a == b for (a, b) in zip (al, bl)])
+
+        eq = lambda a, b : a == b
+
+        def field_is_correct(field, equal_fn, a1, a2):
+            equal = equal_fn(a1, a2)
+            if not equal:
+                print("\nField {} is incorrect: {} != {}".format(field, a1, a2))
+            return equal
+
+        min_ = field_is_correct("min", list_eq, internal.min, target["min"])
+        opt_ = field_is_correct("opt", list_eq, internal.opt, target["opt"])
+        max_ = field_is_correct("max", list_eq, internal.max, target["max"])
+        is_dynamic_ = field_is_correct("is_dynamic", eq, internal.input_is_dynamic, target["input_is_dynamic"])
+        explicit_set_dtype_ = field_is_correct("explicit_dtype", eq, internal._explicit_set_dtype, target["explicit_set_dtype"])
+        dtype_ = field_is_correct("dtype", eq, int(internal.dtype), int(target["dtype"]))
+        format_ = field_is_correct("format", eq, int(internal.format), int(target["format"]))
+
+        return all([min_,opt_,max_,is_dynamic_,explicit_set_dtype_,dtype_,format_])
+
+
+    def test_infer_from_example_tensor(self):
+        shape = [1, 3, 255, 255]
+        target = {
+            "min": shape,
+            "opt": shape,
+            "max": shape,
+            "input_is_dynamic": False,
+            "dtype": torchtrt.dtype.half,
+            "format": torchtrt.TensorFormat.contiguous,
+            "explicit_set_dtype": True
+        }
+
+        example_tensor = torch.randn(shape).half()
+        i = torchtrt.Input._from_tensor(example_tensor)
+        self.assertTrue(self._verify_correctness(i, target))
+
+
+    def test_static_shape(self):
+        shape = [1, 3, 255, 255]
+        target = {
+            "min": shape,
+            "opt": shape,
+            "max": shape,
+            "input_is_dynamic": False,
+            "dtype": torchtrt.dtype.unknown,
+            "format": torchtrt.TensorFormat.contiguous,
+            "explicit_set_dtype": False
+        }
+
+        i = torchtrt.Input(shape)
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(tuple(shape))
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(torch.randn(shape).shape)
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(shape=shape)
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(shape=tuple(shape))
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(shape=torch.randn(shape).shape)
+        self.assertTrue(self._verify_correctness(i, target))
+
+    def test_data_type(self):
+        shape = [1, 3, 255, 255]
+        target = {
+            "min": shape,
+            "opt": shape,
+            "max": shape,
+            "input_is_dynamic": False,
+            "dtype": torchtrt.dtype.half,
+            "format": torchtrt.TensorFormat.contiguous,
+            "explicit_set_dtype": True
+        }
+
+        i = torchtrt.Input(shape, dtype=torchtrt.dtype.half)
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(shape, dtype=torch.half)
+        self.assertTrue(self._verify_correctness(i, target))
+
+    def test_tensor_format(self):
+        shape = [1, 3, 255, 255]
+        target = {
+            "min": shape,
+            "opt": shape,
+            "max": shape,
+            "input_is_dynamic": False,
+            "dtype": torchtrt.dtype.unknown,
+            "format": torchtrt.TensorFormat.channels_last,
+            "explicit_set_dtype": False
+        }
+
+        i = torchtrt.Input(shape, format=torchtrt.TensorFormat.channels_last)
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(shape, format=torch.channels_last)
+        self.assertTrue(self._verify_correctness(i, target))
+
+    def test_dynamic_shape(self):
+        min_shape = [1, 3, 128, 128]
+        opt_shape = [1, 3, 256, 256]
+        max_shape = [1, 3, 512, 512]
+        target = {
+            "min": min_shape,
+            "opt": opt_shape,
+            "max": max_shape,
+            "input_is_dynamic": True,
+            "dtype": torchtrt.dtype.unknown,
+            "format": torchtrt.TensorFormat.contiguous,
+            "explicit_set_dtype": False
+        }
+
+        i = torchtrt.Input(min_shape=min_shape, opt_shape=opt_shape, max_shape=max_shape)
+        self.assertTrue(self._verify_correctness(i, target))
+
+        i = torchtrt.Input(min_shape=tuple(min_shape), opt_shape=tuple(opt_shape), max_shape=tuple(max_shape))
+        self.assertTrue(self._verify_correctness(i, target))
+
+        tensor_shape = lambda shape: torch.randn(shape).shape
+        i = torchtrt.Input(min_shape=tensor_shape(min_shape), opt_shape=tensor_shape(opt_shape), max_shape=tensor_shape(max_shape))
+        self.assertTrue(self._verify_correctness(i, target))
 
 def test_suite():
     suite = unittest.TestSuite()
@@ -371,6 +504,7 @@ def test_suite():
         TestModuleFallbackToTorch.parametrize(TestModuleFallbackToTorch, model=models.resnet18(pretrained=True)))
     suite.addTest(unittest.makeSuite(TestCheckMethodOpSupport))
     suite.addTest(unittest.makeSuite(TestDevice))
+    suite.addTest(unittest.makeSuite(TestInput))
 
     return suite