Repository: pytorch/TensorRT

Issue:
✨[Feature] Allow specification of integer range for `Input` class
Prototype + 1/2 MVP of RFC #1425 

See RFC for details.


Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/_Input.py b/py/torch_tensorrt/_Input.py
--- a/py/torch_tensorrt/_Input.py
+++ b/py/torch_tensorrt/_Input.py
@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Tuple, Optional
 
 import torch
 
@@ -38,6 +38,10 @@ class _ShapeMode(Enum):
         _enums.TensorFormat.contiguous
     )  #: The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
 
+    DOMAIN_OFFSET = 2.0
+    low_tensor_domain_incl = 0.0
+    high_tensor_domain_excl = low_tensor_domain_incl + DOMAIN_OFFSET
+
     def __init__(self, *args, **kwargs):
         """__init__ Method for torch_tensorrt.Input
 
@@ -56,6 +60,8 @@ def __init__(self, *args, **kwargs):
                 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 torch_tensorrt.dtype): Expected data type for input tensor (default: torch_tensorrt.dtype.float32)
             format (torch.memory_format or torch_tensorrt.TensorFormat): The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
+            tensor_domain (Tuple(float, float), optional): The domain of allowed values for the tensor, as interval notation: [tensor_domain[0], tensor_domain[1]).
+                Note: Entering "None" (or not specifying) will set the bound to [0, 2)
 
         Examples:
             - Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)
@@ -138,18 +144,31 @@ def __init__(self, *args, **kwargs):
         if "format" in kwargs:
             self.format = Input._parse_format(kwargs["format"])
 
+        if "tensor_domain" in kwargs:
+            domain = kwargs["tensor_domain"]
+        else:
+            domain = None
+
+        self.tensor_domain = Input._parse_tensor_domain(domain)
+
     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)
+            return "Input(shape={}, dtype={}, format={}, domain=[{}, {}))".format(
+                self.shape,
+                str(self.dtype),
+                str(self.format),
+                str(self.tensor_domain[0]),
+                str(self.tensor_domain[1]),
             )
         elif self.shape_mode == Input._ShapeMode.DYNAMIC:
-            return "Input(min_shape={}, opt_shape={}, max_shape={}, dtype={}, format={})".format(
+            return "Input(min_shape={}, opt_shape={}, max_shape={}, dtype={}, format={}, domain=[{}, {}))".format(
                 self.shape["min_shape"],
                 self.shape["opt_shape"],
                 self.shape["max_shape"],
                 str(self.dtype),
                 str(self.format),
+                str(self.tensor_domain[0]),
+                str(self.tensor_domain[1]),
             )
         else:
             raise RuntimeError("Unknown input shape mode")
@@ -203,6 +222,8 @@ def _to_internal(self) -> _C.Input:
         internal_in.dtype = Input._parse_dtype(self.dtype)
         internal_in._explicit_set_dtype = self._explicit_set_dtype
         internal_in.format = Input._parse_format(self.format)
+
+        internal_in.tensor_domain = Input._parse_tensor_domain(self.tensor_domain)
         return internal_in
 
     @staticmethod
@@ -267,6 +288,51 @@ def _parse_format(format: Any) -> _enums.TensorFormat:
                 "Tensor format needs to be specified with either torch.memory_format or torch_tensorrt.TensorFormat"
             )
 
+    @staticmethod
+    def _parse_tensor_domain(domain: Optional[Tuple[float, float]]) -> Tuple:
+        """
+        Produce a tuple of integers which specifies a tensor domain in the interval format: [lo, hi)
+
+        Args:
+            domain (Tuple[int, int]): A tuple of integers (or NoneTypes) to verify
+
+        Returns:
+            A tuple of two int32_t-valid integers
+        """
+        if domain is None:
+            result_domain = (
+                Input.low_tensor_domain_incl,
+                Input.high_tensor_domain_excl,
+            )
+        elif len(domain) == 2:
+            domain_lo, domain_hi = domain
+
+            # Validate type and provided values for domain
+            valid_type_lo = isinstance(domain_lo, int) or isinstance(domain_lo, float)
+            valid_type_hi = isinstance(domain_hi, int) or isinstance(domain_hi, float)
+
+            if not valid_type_lo:
+                raise ValueError(
+                    f"Expected value for tensor domain low specifier, got {domain_lo}"
+                )
+            elif not valid_type_hi:
+                raise ValueError(
+                    f"Expected value for tensor domain high specifier, got {domain_hi}"
+                )
+
+            if domain_hi <= domain_lo:
+                raise ValueError(
+                    "Expected provided integer range to have low tensor domain value "
+                    + f"< high tensor domain value, got invalid range [{domain_lo}, {domain_hi})"
+                )
+            result_domain = (float(domain_lo), float(domain_hi))
+        else:
+            raise ValueError(
+                f"Expected 2 values for domain, got {len(domain)}: {domain}"
+            )
+
+        return result_domain
+
     @classmethod
     def from_tensor(cls, t: torch.Tensor) -> "Input":
         """
diff --git a/py/torch_tensorrt/ts/_compile_spec.py b/py/torch_tensorrt/ts/_compile_spec.py
--- a/py/torch_tensorrt/ts/_compile_spec.py
+++ b/py/torch_tensorrt/ts/_compile_spec.py
@@ -17,6 +17,7 @@ def _internal_input_to_torch_class_input(i: _C.Input) -> torch.classes.tensorrt.
     clone._set_opt(i.opt)
     clone._set_max(i.max)
     clone._set_dtype(i.dtype)
+    clone._set_tensor_domain(i.tensor_domain)
     clone._set_format(i.format)
     clone._set_input_is_dynamic(i.input_is_dynamic)
     clone._set_explicit_set_dtype(i._explicit_set_dtype)


Test patch:
diff --git a/tests/cpp/test_collections.cpp b/tests/cpp/test_collections.cpp
--- a/tests/cpp/test_collections.cpp
+++ b/tests/cpp/test_collections.cpp
@@ -89,6 +89,51 @@ TEST(CppAPITests, TestCollectionStandardTensorInputLongDtype) {
       out.toTensor().to(torch::kFloat), trt_out.toTensor().to(torch::kFloat)));
 }
 
+TEST(CppAPITests, TestSpecifyDomainStandardTensorInput) {
+  std::string path = "tests/modules/tuple_input_scripted.jit.pt";
+  torch::Tensor in0 = torch::randn({1, 3, 512, 512}, torch::kCUDA).to(torch::kHalf);
+
+  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> complex_inputs, complex_inputs_list;
+  std::tuple<torch::jit::IValue, torch::jit::IValue> input_tuple(in0, in0);
+
+  complex_inputs.push_back(input_tuple);
+
+  auto out = mod.forward(complex_inputs);
+
+  // Specify input tensor domain argument
+  auto tensor_domain = std::vector<double>{35, 377};
+  auto input_shape = torch_tensorrt::Input(in0.sizes(), torch_tensorrt::DataType::kHalf, tensor_domain);
+
+  auto input_shape_ivalue = torch::jit::IValue(std::move(c10::make_intrusive<torch_tensorrt::Input>(input_shape)));
+
+  std::tuple<torch::jit::IValue, torch::jit::IValue> input_shape_tuple(input_shape_ivalue, input_shape_ivalue);
+
+  torch::jit::IValue complex_input_shape(input_shape_tuple);
+  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.min_block_size = 1;
+
+  // // FP16 execution
+  compile_settings.enabled_precisions = {torch::kHalf};
+  // // Compile module
+  auto trt_mod = torch_tensorrt::torchscript::compile(mod, compile_settings);
+  auto trt_out = trt_mod.forward(complex_inputs);
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(out.toTensor(), trt_out.toTensor()));
+}
+
 TEST(CppAPITests, TestCollectionTupleInput) {
   std::string path = "tests/modules/tuple_input_scripted.jit.pt";
   torch::Tensor in0 = torch::randn({1, 3, 512, 512}, torch::kCUDA).to(torch::kHalf);
diff --git a/tests/py/api/test_collections.py b/tests/py/api/test_collections.py
--- a/tests/py/api/test_collections.py
+++ b/tests/py/api/test_collections.py
@@ -80,6 +80,35 @@ def test_compile(self):
         )
 
 
+class TestStandardTensorInputDomain(unittest.TestCase):
+    def test_compile(self):
+
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+        self.model = (
+            torch.jit.load(MODULE_DIR + "/standard_tensor_input_scripted.jit.pt")
+            .eval()
+            .to("cuda")
+        )
+
+        compile_spec = {
+            "inputs": [
+                torchtrt.Input(self.input.shape, tensor_domain=(70.8, 800)),
+                torchtrt.Input(self.input.shape, tensor_domain=(-20, -17.9)),
+            ],
+            "device": torchtrt.Device("gpu:0"),
+            "enabled_precisions": {torch.float},
+        }
+
+        trt_mod = torchtrt.ts.compile(self.model, **compile_spec)
+        cos_sim = cosine_similarity(
+            self.model(self.input, self.input), trt_mod(self.input, self.input)
+        )
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"standard_tensor_input_scripted with tensor domain specified TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
+
+
 class TestTupleInput(unittest.TestCase):
     def test_compile(self):
 
diff --git a/tests/py/api/test_operator_fallback.py b/tests/py/api/test_operator_fallback.py
--- a/tests/py/api/test_operator_fallback.py
+++ b/tests/py/api/test_operator_fallback.py
@@ -31,6 +31,32 @@ def test_fallback_resnet18(self):
             msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
         )
 
+    def test_fallback_resnet18_with_tensor_domain(self):
+        self.model = models.resnet18(pretrained=True).eval().to("cuda")
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+        compile_spec = {
+            "inputs": [
+                torchtrt.Input(
+                    self.input.shape,
+                    dtype=torch.float,
+                    format=torch.contiguous_format,
+                    tensor_domain=(-0.5, 0.5),
+                )
+            ],
+            "device": {
+                "device_type": torchtrt.DeviceType.GPU,
+                "gpu_id": 0,
+            },
+            "enabled_precisions": {torch.float},
+            "torch_executed_ops": ["aten::add"],
+        }
+        trt_mod = torchtrt.compile(self.model, **compile_spec)
+        cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
+
     def test_fallback_mobilenet_v2(self):
         self.model = models.mobilenet_v2(pretrained=True).eval().to("cuda")
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")