❓ [Question] How can I move the converted tensorRT model in a Jetson system?

[mjack3]

❓ Question

I optimized a pytorch module with torch-TensorRT. How can I move the engine to a Jetson?

What you have already tried

I tried torch.jit.load('trt_traced_model.ts')

but get torch.torch.classes.tensorrt.Engine error

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• PyTorch Version (e.g., 1.0): 10.0
• OS (e.g., Linux): ARM Ubuntu 18
• How you installed PyTorch : pip from offical Nvidia support
• Python version: 3.6
• CUDA version: 10.2
• GPU models and configuration: Jetson NX

Additional context

I have a Jetson NX system with jetpack 4.6, torch v0.10.0 and torchvision v0.11.0 where I want to deploy a tensorRT model.

For that in my main computer I installed this repository and converted my model to tensorRT successfully. I need to move it into the Jetson for production.

This is the code that I use to export to tensorRT (main computer)

model.cuda().eval()
model = torch.jit.trace(model, [torch.rand(1, 3, 224, 224).cuda()])
trt_model_fp32 = torch_tensorrt.compile(model,
                                            inputs=[torch_tensorrt.Input((1, 3, 224, 224))],
                                            enabled_precisions=torch.float32,  # Run with FP32
                                            )
torch.jit.save(trt_model_fp32, dir)

This is in my Jetson

model = torch.jit.load(dir)

but i get torch.torch.classes.tensorrt.Engine error

Jetson hasn't installed torch-tensorRT. How can I move the tensorRT model? Do I need to install this repo also in the Jetson?

Thanks!

[narendasan]

(moving this to a discussion so people can find it later)

TensorRT engines must be built on the device you intend to deploy it on. So the workflow I recommend is you export your model to TorchScript on your host machine, then move that file to the Jetson where you would do the Torch-TensorRT compilation.

Basically just separate the code you provided into two scripts, one run on the host and one run on the Jetson.

host.py

model.cuda().eval()
model = torch.jit.trace(model, [torch.rand(1, 3, 224, 224).cuda()])
model.save("my_model.ts")

device_compile.py

model = torch.jit.load("my_model.ts")
trt_model_fp32 = torch_tensorrt.compile(model,
                                        inputs=[torch_tensorrt.Input((1, 3, 224, 224))],
                                        enabled_precisions=torch.float32)
torch.jit.save(trt_model_fp32, dir)

[mjack3]

Wow! Thank you for your answer!

Could you respond one last?

What's the way to install torch-TensorRT in the Jetson device? Using pip give me an error. Does it need be from source?

[narendasan]

Yeah for now it needs to be from source. Right now Jetson support is in a weird spot due to a recent Jetpack upgrade. If you are on Jetpack 4.6, building the v1.0.0 tag should work out of the box (https://nvidia.github.io/Torch-TensorRT/tutorials/installation.html#building-natively-on-aarch64-jetson). If you are using 5.0DP, we are going to add support very soon.

[mjack3]

Thanks for your help. I installed torch_tensorRT on Jetson device sucessfully. I think that the documentation for installing torch_tensorrt in Jetson is confusing so I will try to sum up the steps needed here:

1. Install Go
2. Install Bazel (from source or from bazelisk). I tested both and the second one is easier and faster.

If you see something with go version and bazel --version you are good

4. Follow the steps of Building natively on aarch64.

Notes:

1 - Just TensorRT 1.1.0 is compatible with JetPack < 5
2 - python3 setup.py install --use-cxx11-abi will install torch_tensorrt in /usr/local/lib/python3.6/dist-packages but if you have torch installed on $HOME/.local/lib/python3.6/site-packages/torch you need to add --user. This also solved a problem of permissions to create the final torch_tensorrt folder

[narendasan]

You just need to install the bazel binary not necessarily bazelisk. At least that is what we do. Yes right now we do not support Jetpack 5.0, that will be coming in a point release soon.