JetsonでTensorRTのセットアップとEngine化を行う手順まとめ (4)

Jetson で TensorRT を使うとき、最初に欲しいのは「とにかく .onnx を .engine にして動かす」流れです。

この記事では、TensorRT の Python 環境構築から、trtexec を使った engine 化、INT8 / FP16 の考え方、Engine 情報の確認、Python サンプルまでまとめます。

TensorRT Python 環境を整える

まず Python から TensorRT を扱える状態にします。

sudo apt update
sudo apt install python3-libnvinfer python3-libnvinfer-dev python3-libnvinfer-plugin

仮想環境（venv）から使う場合は、system site packages を参照できるようにしておくと扱いやすいです。

vi python/venv/pyvenv.cfg

include-system-site-packages = true

確認:

python3 -c "import tensorrt; print(tensorrt.version)"

PyCUDA も入れておくとサンプル実装が書きやすくなります。

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export CPATH=/usr/local/cuda/include:$CPATH

pip3 install pycuda

trtexec で ONNX を Engine 化する

基本コマンドは次のような形です。

/usr/src/tensorrt/bin/trtexec --onnx=model.onnx --saveEngine=model.engine --fp16 --useCudaGraph

よく使うオプション:

• --builderOptimizationLevel=5
• --useCudaGraph
• --explicitBatch
• --minShapes --optShapes --maxShapes
• --workspace=N
• --verbose
• --fp16
• --separateProfileRun

推奨例:

/usr/src/tensorrt/bin/trtexec \
  --onnx=model.onnx \
  --saveEngine=model.engine \
  --fp16 \
  --useCudaGraph \
  --builderOptimizationLevel=3 \
  --avgTiming=16 \
  --separateProfileRun \
  --verbose

INT8 量子化の例

INT8 化まで踏み込みたい場合は、キャリブレーション素材を用意して次のような形になります。

/usr/src/tensorrt/bin/trtexec \
  --onnx=model.onnx \
  --saveEngine=model_int8.engine \
  --int8 \
  --calib=calib.cache \
  --fp16 \
  --useCudaGraph \
  --avgTiming=16 \
  --separateProfileRun \
  --sparsity=enable \
  --memPoolSize=workspace:4096 \
  --verbose

ここでのポイントは、INT8 は便利ですがキャリブレーション素材が必要で、すぐに動作確認したいだけなら FP16 の方が入りやすいことが多い、という点です。

ベンチマークする

生成した Engine は trtexec でそのままベンチマークできます。

/usr/src/tensorrt/bin/trtexec --loadEngine=model.engine

まずはここで、「モデルは変換できたのか」「速度はどのくらいか」を切り分けるのが楽です。

Engine の入出力を確認する

Engine の中身を確認したいときは、Python からロードして入出力テンソルを列挙できます。

import tensorrt as trt


def inspect_engine(engine_path):
    logger = trt.Logger(trt.Logger.WARNING)
    trt.init_libnvinfer_plugins(logger, "")
    with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())

    print(f"[INFO] Engine: {engine_path}")
    print(f"[INFO] num_io_tensors = {engine.num_io_tensors}")

    for i in range(engine.num_io_tensors):
        name = engine.get_tensor_name(i)
        dtype = engine.get_tensor_dtype(name)
        shape = engine.get_tensor_shape(name)
        mode = engine.get_tensor_mode(name)
        print(f" - {mode.name}: {name} | shape={shape} | dtype={dtype}")

shape や dtype の理解がずれていると、その後の Python / C++ 実装で詰まりやすいので、ここはかなり重要です。

Python から Engine を叩くサンプル

実行の最小例としては、次のような形になります。

import numpy as np
import tensorrt as trt
import pycuda.driver as cuda

cuda.init()
device = cuda.Device(0)
ctx = device.make_context()

class TrtRunner:
    def __init__(self, engine_path):
        logger = trt.Logger(trt.Logger.WARNING)
        trt.init_libnvinfer_plugins(logger, "")
        with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self.stream = cuda.Stream()

元メモにはもう少し長いサンプルがありますが、記事としては「Engine をロードし、入力 shape を設定し、非同期実行して結果を戻す」という流れが分かれば十分だと思います。

まとめ

Jetson で TensorRT を使い始めるときは、まず次の流れを押さえておくと進めやすいです。

1. TensorRT Python 環境を整える
2. trtexec で .onnx を .engine に変換する
3. FP16 / INT8 の方針を決める
4. Engine をベンチマークする
5. 入出力テンソルを確認する
6. Python か C++ から実行してみる

まずは trtexec で素早く動かし、その後にコードへ落としていく方が失敗しにくいです。

次は、TensorRT / CUDA を C++ で扱うときの実装と最適化ポイントを整理します。