Repository: pytorch/TensorRT

Issue:
🐛 [Bug] An wrong function call for `torch._C._jit_to_tensorrt` on `spec` in tutorial code.
##  Bug Description

This is a bug in a tutorial below.

* Web: https://nvidia.github.io/TRTorch/tutorials/use_from_pytorch.html
* Original ReST doc: https://github.com/NVIDIA/TRTorch/blob/51a2043217a3f6c93f393169961750733c0d26ec/docsrc/tutorials/use_from_pytorch.rst

To compile a model, `torch._C._jit_to_tensorrt()` is called with two parameters, `script_model._c` and `spec`.
But, a following error happened.

```
Traceback (most recent call last):
  File "simple.py", line 34, in <module>
    main()
  File "simple.py", line 27, in main
    trt_model = torch._C._jit_to_tensorrt(script_model._c, spec)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<string>", line 6, in __setstate__
                self.__processed_module = state[1]
                self.__create_backend()
                self.__handles = self.__backend.compile(self.__processed_module, self.__method_compile_spec)
                                 ~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

RuntimeError: isGenericDict() INTERNAL ASSERT FAILED at "/opt/python/cp36-cp36m/lib/python3.6/site-packages/torch/include/ATen/core/ivalue_inl.h":830, please report a bug to PyTorch. Expected GenericDict but got Object
```

This error looks like caused by wrong parameter. In test code (https://github.com/NVIDIA/TRTorch/blob/b93627ecc7bb7123d2e32aff9d762dd0e6bc3166/tests/py/test_to_backend_api.py), this API is called with a different parameter like below. It looks like making additional `dict` with `forward` key is necessary.

```
trt_model = torch._C._jit_to_tensorrt(script_model._c, {'forward': spec})
```

## To Reproduce

Steps to reproduce the behavior:

1. Launch NGC TRT container, `nvcr.io/nvidia/tensorrt:20.03-py3`.
1. Install required libraries
    - `pip install https://github.com/NVIDIA/TRTorch/releases/download/v0.1.0/trtorch-0.1.0-cp36-cp36m-linux_x86_64.whl torchvision==0.7.0`
1. Run each step described in the tutorial.

My repro code is below. Note that I modified the tutorial code to use FP32 for test.

```
import torch
import trtorch
import torchvision.models as models

def main():
    model = models.mobilenet_v2(pretrained=True)
    script_model = torch.jit.script(model)

    spec = {
        "forward": trtorch.TensorRTCompileSpec({
            "input_shapes": [[1, 3, 300, 300]],
            "op_precision": torch.float32,
            "refit": False,
            "debug": False,
            "strict_types": False,
            "allow_gpu_fallback": True,
            "device_type": "gpu",
            "capability": trtorch.EngineCapability.default,
            "num_min_timing_iters": 2,
            "num_avg_timing_iters": 1,
            "max_batch_size": 0,
        })
    }
    # trt_model = torch._C._jit_to_tensorrt(script_model._c, {'forward': spec})
    trt_model = torch._C._jit_to_tensorrt(script_model._c, spec)

    x = torch.randn((1, 3, 300, 300)).to('cuda').to(torch.float32)
    print(trt_model.forward(x))


if __name__ == '__main__':
    main()
```

## Expected behavior

No exception.

## Environment

> Build information about the TRTorch compiler can be found by turning on debug messages

 - PyTorch Version (e.g., 1.0): 1.6.0
 - CPU Architecture: Intel(R) Core(TM) i7-6850K
 - OS (e.g., Linux): Ubuntu 18.04.5 LTS (host) / Ubuntu 18.04.4 LTS (container)
 - How you installed PyTorch (`conda`, `pip`, `libtorch`, source): pip
 - Build command you used (if compiling from source): N/A
 - Are you using local sources or building from archives: N/A
 - Python version: 3.6.9
 - CUDA version: 10.2.89
 - GPU models and configuration: TitanX (Pascal)
 - Any other relevant information: N/A

## Additional context

N/A



Discussion:
It seems like to me there might be some confusing variable names. `spec` in the test is of type `trtorch.TensorRTCompileSpec`, `torch._C._jit_to_tensorrt` takes a `torch._C.ScriptModule` and a Dictionary `{ str : torch.TensorRTCompileSpec }`. So in the test we need to make a dictionary `{'forward': spec}`. Hence 

```py
trt_model = torch._C._jit_to_tensorrt(script_model._c, {'forward': spec})
```

`spec` in the docs unless I am misreading it is itself a Dictionary `{ str : trtorch.TensorRTCompileSpec }` so we should be able to call it with 

```py
trt_model = torch._C._jit_to_tensorrt(script_model._c, spec)
```

 as shown in the documentation (https://nvidia.github.io/TRTorch/tutorials/use_from_pytorch.html).

It would probably be a good idea to change the name of one of these to avoid confusion.
Let me add one comment about `spec` in the test code.

It is defined like below therefore current `torch._C._jit_to_tensorrt` seems to take a Dictionary `{ str : { str : torch.TensorRTCompileSpec }}` instead of simple `{ str : torch.TensorRTCompileSpec }`.
https://github.com/NVIDIA/TRTorch/blob/b93627ecc7bb7123d2e32aff9d762dd0e6bc3166/tests/py/test_to_backend_api.py#L14-L29

Is that expected?

Resolving patch:
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
@@ -147,10 +147,6 @@ def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
         assert isinstance(compile_spec["strict_types"], bool)
         info.strict_types = compile_spec["strict_types"]
 
-    if "allow_gpu_fallback" in compile_spec:
-        assert isinstance(compile_spec["allow_gpu_fallback"], bool)
-        info.allow_gpu_fallback = compile_spec["allow_gpu_fallback"]
-
     if "device" in compile_spec:
         info.device = _parse_device(compile_spec["device"])
 
@@ -177,7 +173,7 @@ def _parse_compile_spec(compile_spec: Dict[str, Any]) -> trtorch._C.CompileSpec:
     return info
 
 
-def TensorRTCompileSpec(compile_spec: Dict[str, Any]):
+def TensorRTCompileSpec(compile_spec: Dict[str, Any]) -> torch.classes.tensorrt.CompileSpec:
     """
     Utility to create a formated spec dictionary for using the PyTorch TensorRT backend
 
@@ -199,10 +195,10 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]):
                             } # 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)
-                        "gpu_id": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)
-                        "dla_core": 0, # (DLA only) Target dla core id to run engine
-                        "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
+                            "device_type": torch.device("cuda"), # Type of device to run engine on (for DLA use trtorch.DeviceType.DLA)
+                            "gpu_id": 0, # Target gpu id to run engine (Use Xavier as gpu id for DLA)
+                            "dla_core": 0, # (DLA only) Target dla core id to run engine
+                            "allow_gpu_fallback": false, # (DLA only) Allow layers unsupported on DLA to run on GPU
                         },
                         "op_precision": torch.half, # Operating precision set to FP16
                         "refit": False, # enable refit
@@ -235,14 +231,13 @@ def TensorRTCompileSpec(compile_spec: Dict[str, Any]):
         ir.set_max(i.max)
         backend_spec.append_input_range(ir)
 
-    for i in parsed_spec.device:
-        ir = torch.classes.tensorrt.Device()
-        ir.set_device_type(i.device_type)
-        ir.set_gpu_id(i.gpu_id)
-        ir.set_dla_core(i.dla_core)
-        ir.set_allow_gpu_fallback(i.allow_gpu_fallback)
-        backend_spec.set_device(ir)
+    d = torch.classes.tensorrt.Device()
+    d.set_device_type(int(parsed_spec.device.device_type))
+    d.set_gpu_id(parsed_spec.device.gpu_id)
+    d.set_dla_core(parsed_spec.device.dla_core)
+    d.set_allow_gpu_fallback(parsed_spec.device.allow_gpu_fallback)
 
+    backend_spec.set_device(d)
     backend_spec.set_op_precision(int(parsed_spec.op_precision))
     backend_spec.set_refit(parsed_spec.refit)
     backend_spec.set_debug(parsed_spec.debug)


Test patch:
diff --git a/tests/py/BUILD b/tests/py/BUILD
--- a/tests/py/BUILD
+++ b/tests/py/BUILD
@@ -17,7 +17,8 @@ py_test(
     ] + select({
 	":aarch64_linux": [
 		"test_api_dla.py"
-	]
+	],
+    "//conditions:default" : []
     }),
     deps = [
         requirement("torchvision")
diff --git a/tests/py/test_to_backend_api.py b/tests/py/test_to_backend_api.py
--- a/tests/py/test_to_backend_api.py
+++ b/tests/py/test_to_backend_api.py
@@ -19,8 +19,12 @@ def setUp(self):
                     "refit": False,
                     "debug": False,
                     "strict_types": False,
-                    "allow_gpu_fallback": True,
-                    "device_type": "gpu",
+                    "device": {
+                        "device_type": trtorch.DeviceType.GPU,
+                        "gpu_id": 0,
+                        "dla_core": 0,
+                        "allow_gpu_fallback": True
+                    },
                     "capability": trtorch.EngineCapability.default,
                     "num_min_timing_iters": 2,
                     "num_avg_timing_iters": 1,
@@ -29,14 +33,14 @@ def setUp(self):
         }
 
     def test_to_backend_lowering(self):
-        trt_mod = torch._C._jit_to_tensorrt(self.scripted_model._c, {"forward": self.spec})
+        trt_mod = torch._C._jit_to_backend("tensorrt", self.scripted_model, self.spec)
         same = (trt_mod.forward(self.input) - self.scripted_model(self.input)).abs().max()
         self.assertTrue(same < 2e-3)
 
 
 def test_suite():
     suite = unittest.TestSuite()
-    suite.addTest(TestToBackendLowering.parametrize(TestToBackendLowering, model=models.mobilenet_v2(pretrained=True)))
+    suite.addTest(TestToBackendLowering.parametrize(TestToBackendLowering, model=models.resnet18(pretrained=True)))
 
     return suite