PyTorch→ONNX→TensorRTの変換

研究室でTensorRTのモデルをC++で動かして推論をしているのですが、毎回変換の細かい手順を忘れるので書き残しておきます。
特にONNX→TensorRTの変換に癖があって、私の経験上だと入力テンソルの大きさが可変だとだいたいうまくいきません。

PyTorch→ONNX

convert.py

import torch
import torch.onnx

model = MyModel()
model.load_state_dict(torch.load('hoge.pth'))

torch.onnx.export(model, torch.randn((1,3,720,1280)), 'hoge.onnx')

ONNX→TensorRT

trtexecを使います

trtexec --onnx=hoge.onnx --saveEngine=hoge.trt