Repository: pytorch/TensorRT

Issue:
✨[Feature] Let users set input shape, dtype and format using example tensors
**Is your feature request related to a problem? Please describe.**
<!--A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]-->
Now that we are going to be able to specify tensor properties for inputs such as data type and format, it should be easier to use a tensor I already have as an example of the shape dtype and format I want.

**Describe the solution you'd like**
<!--A clear and concise description of what you want to happen.-->
compile should be able to accept a list of Tensors to represent my graph inputs. From those tensors compile should be able to infer shape, data type and tensor format for the user. These should be logged in the debug log for easy debugging. 

**Describe alternatives you've considered**
<!--A clear and concise description of any alternative solutions or features you've considered.-->
Note: I dont think theres a way to support dynamic shape using this method, so users will need to use `trtorch.Input` for that 

**Additional context**
<!--Add any other context or screenshots about the feature request here.-->



Discussion:
Why, if you supply 3 Tensors for min/opt/max, it's no less flexibility than we have with size arguments. 
This also correlates with https://github.com/NVIDIA/TRTorch/issues/538 

> Why, if you supply 3 Tensors for min/opt/max, it's no less flexibility than we have with size arguments.

There are some issues I see with defining min/opt/max with example tensors, at least how i pictured it. For instance: 

```py
trtorch.compile(my_mod, {"inputs": [example_in]})
```
From `example_in` we can infer shape, dtype and tensor format directly from the Tensor and create a `trtorch.Input` internally

In the case of dynamic shape, we need some sort of grouping for the tensor, so maybe you do something like 

```py
trtorch.compile(my_mod, {"inputs": [
    {"min": example_in_min, 
     "opt": example_in_opt, 
     "max": example_in_max}
]})
``` 
But, in this case, what do we do if something like the dtype or tensor format is inconstant between tensors? 
Also doesnt seem like this is that different than doing at this point: 

```py
trtorch.compile(my_mod, {"inputs": [
    trtorch.Input(min_shape=example_in_min.shape(), 
        opt_shape=example_in_opt.shape(), 
        max_shape=example_in_max.shape(), dtype=torch.float)
]})
``` 
We could maybe have `trtorch.Input` pull out the shape and error out if the types are inconstant. Really though we wanted to give users a bypass around `trtorch.Input` with this feature.








Resolving patch:
diff --git a/py/trtorch/Input.py b/py/trtorch/Input.py
--- a/py/trtorch/Input.py
+++ b/py/trtorch/Input.py
@@ -196,3 +196,16 @@ def _parse_format(format: Any) -> _types.TensorFormat:
         else:
             raise TypeError(
                 "Tensor format needs to be specified with either torch.memory_format or trtorch.TensorFormat")
+
+    @classmethod
+    def _from_tensor(cls, t: torch.Tensor):
+        if not any([
+                t.is_contiguous(memory_format=torch.contiguous_format),
+                t.is_contiguous(memory_format=torch.channels_last)
+        ]):
+            raise ValueError(
+                "Tensor does not have a supported contiguous memory format, supported formats are contiguous or channel_last"
+            )
+        frmt = torch.contiguous_format if t.is_contiguous(
+            memory_format=torch.contiguous_format) else torch.channels_last
+        return cls(shape=t.shape, dtype=t.dtype, format=frmt)
diff --git a/py/trtorch/_compile_spec.py b/py/trtorch/_compile_spec.py
--- a/py/trtorch/_compile_spec.py
+++ b/py/trtorch/_compile_spec.py
@@ -174,7 +174,12 @@ def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
         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 not all([isinstance(i, torch.Tensor) or isinstance(i, trtorch.Input) for i in compile_spec["inputs"]]):
+            raise KeyError("Input specs should be either trtorch.Input or torch.Tensor, found types: {}".format(
+                [typeof(i) for i in compile_spec["inputs"]]))
+
+        inputs = [trtorch.Input._from_tensor(i) if isinstance(i, torch.Tensor) else i for i in compile_spec["inputs"]]
+        info.inputs = [i._to_internal() for i in inputs]
 
     if "op_precision" in compile_spec and "enabled_precisions" in compile_spec:
         raise KeyError(


Test patch:
diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD
--- a/tests/cpp/BUILD
+++ b/tests/cpp/BUILD
@@ -16,7 +16,8 @@ test_suite(
         ":test_modules_as_engines",
         ":test_multiple_registered_engines",
         ":test_serialization",
-        ":test_module_fallback"
+        ":test_module_fallback",
+        ":test_example_tensors"
     ],
 )
 
@@ -28,7 +29,8 @@ test_suite(
         ":test_modules_as_engines",
         ":test_multiple_registered_engines",
         ":test_serialization",
-        ":test_module_fallback"
+        ":test_module_fallback",
+        ":test_example_tensors"
     ],
 )
 
@@ -43,6 +45,17 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "test_example_tensors",
+    srcs = ["test_example_tensors.cpp"],
+    data = [
+        "//tests/modules:jit_models",
+    ],
+    deps = [
+        ":cpp_api_test",
+    ],
+)
+
 cc_test(
     name = "test_serialization",
     srcs = ["test_serialization.cpp"],
diff --git a/tests/cpp/test_example_tensors.cpp b/tests/cpp/test_example_tensors.cpp
new file mode 100644
--- /dev/null
+++ b/tests/cpp/test_example_tensors.cpp
@@ -0,0 +1,23 @@
+#include "cpp_api_test.h"
+
+TEST_P(CppAPITests, InputsFromTensors) {
+  std::vector<torch::jit::IValue> jit_inputs_ivalues;
+  std::vector<torch::jit::IValue> trt_inputs_ivalues;
+  for (auto in_shape : input_shapes) {
+    auto in = at::randn(in_shape, {at::kCUDA});
+    jit_inputs_ivalues.push_back(in.clone());
+    trt_inputs_ivalues.push_back(in.clone());
+  }
+
+  auto spec = trtorch::CompileSpec({trt_inputs_ivalues[0].toTensor()});
+
+  auto trt_mod = trtorch::CompileGraph(mod, spec);
+  torch::jit::IValue trt_results_ivalues = trtorch::tests::util::RunModuleForward(trt_mod, trt_inputs_ivalues);
+  std::vector<at::Tensor> trt_results;
+  trt_results.push_back(trt_results_ivalues.toTensor());
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CompiledModuleForwardIsCloseSuite,
+    CppAPITests,
+    testing::Values(PathAndInSize({"tests/modules/resnet18_traced.jit.pt", {{1, 3, 224, 224}}, 2e-5})));
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
@@ -73,6 +73,27 @@ def test_compile_script(self):
         same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
         self.assertTrue(same < 2e-2)
 
+    def test_from_torch_tensor(self):
+        compile_spec = {
+            "inputs": [self.input],
+            "device": {
+                "device_type": trtorch.DeviceType.GPU,
+                "gpu_id": 0,
+            },
+            "enabled_precisions": {torch.float}
+        }
+
+        trt_mod = trtorch.compile(self.scripted_model, compile_spec)
+        same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
+        self.assertTrue(same < 2e-2)
+
+    def test_device(self):
+        compile_spec = {"inputs": [self.input], "device": trtorch.Device("gpu:0"), "enabled_precisions": {torch.float}}
+
+        trt_mod = trtorch.compile(self.scripted_model, compile_spec)
+        same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
+        self.assertTrue(same < 2e-2)
+
 
 class TestCompileHalf(ModelTestCase):