JetsonでTensorRTを使う前に整理したいモデル変換とONNXの話 (2)

Jetson で TensorRT を使うとき、最初に詰まりやすいのは推論コードそのものよりも モデル形式の整理 です。

TensorRT では基本的に .onnx を入力にすることが多く、TensorFlow / Keras / TFLite 系のモデルをどう ONNX に変換するかが最初の山になります。

この記事では、Jetson で TensorRT を使う前提として、モデル変換周りをまとめて整理します。

TensorRT を使う主な方法

TensorRT を使う方法は大きく分けて次の3つです。

• ONNX Runtime + TensorRT Execution Provider
• ONNX Runtime 単体
• TensorRT 単体

それぞれ一長一短がありますが、Jetson で実運用寄りに詰めていくなら、最終的には TensorRT 単体か、少なくとも .engine を意識した構成に寄っていくことが多い印象です。

Jetson での前提

Jetson で TensorRT を使う場合、まず JetPack 側の導入状態を確認します。

sudo apt update
sudo apt install nvidia-jetpack
sudo apt show nvidia-jetpack

Python 側の最低限の準備としては、次のようなモジュールを入れておくと扱いやすいです。

pip install tensorflow
pip install onnxruntime

ここでの onnxruntime は CPU 確認用として使うことがあります。

Jetson 用 onnxruntime-gpu の導入

Jetson で GPU を使う場合は、CPU 版 onnxruntime をそのまま使うのではなく、GPU 対応版を使う方が自然です。

まず CPU 版を削除します。

pip uninstall onnxruntime

その後、aarch64 向けの GPU 版を入れます。

wget https://github.com/ultralytics/assets/releases/download/v0.0.0/onnxruntime_gpu-1.20.0-cp310-cp310-linux_aarch64.whl
pip install onnxruntime_gpu-1.20.0-cp310-cp310-linux_aarch64.whl

TFLite から ONNX へ変換する

tflite から onnx へ変換するには tf2onnx を使います。

pip install tensorflow
pip install -U tf2onnx
pip install onnxruntime

最新版を GitHub から入れる場合は次のようにもできます。

pip install git+https://github.com/onnx/tensorflow-onnx

変換コマンドの基本形は次のとおりです。

python3 -m tf2onnx.convert --tflite ./model.tflite --output ./model.onnx --opset 13 --verbose

JetPack SDK 6.2 周辺では、--opset 13 が安定しやすいケースがありました。

TFL_RELU_0_TO_1 エラー対策

変換後の ONNX モデルに TFL_RELU_0_TO_1 が残っていて、そのままだと扱いにくいケースがあります。

その場合は Clip ノードへ置換する方法があります。

import onnx
from onnx import helper

input_path = "decoder.onnx"
output_path = "decoder_fixed.onnx"

model = onnx.load(input_path)
graph = model.graph

new_nodes = []
for node in graph.node:
    if node.op_type == "TFL_RELU_0_TO_1":
        print(f"置換対象ノード: {node.name}")

        input_name = node.input[0]
        output_name = node.output[0]

        clip_node = helper.make_node(
            "Clip",
            inputs=[input_name],
            outputs=[output_name],
            name=node.name + "_clip",
            min=0.0,
            max=1.0,
        )
        new_nodes.append(clip_node)
    else:
        new_nodes.append(node)

graph.ClearField("node")
graph.node.extend(new_nodes)
onnx.save(model, output_path)
print(f"修正済みモデルを保存しました: {output_path}")

このあたりは「変換が終わった = そのまま TensorRT に食わせられる」とは限らないことの典型例です。

ONNX モデルを軽量化する

変換した ONNX モデルは、そのままでは冗長なことがあります。

onnxsim で簡略化する

pip install onnxsim
python3 -m onnxsim model.onnx model_sim.onnx

onnxruntime-tools で最適化する

pip install onnxruntime-tools

from onnxruntime_tools import optimizer

optimized = optimizer.optimize_model("decoder11_sim.onnx")
optimized.save_model_to_file("decoder11_opt.onnx")

ONNX の入出力 shape を確認する

モデル変換後は、入出力 shape を見ておくと後でかなり助かります。

import onnx

model = onnx.load("decoder.onnx")

print("=== Inputs ===")
for input in model.graph.input:
    dims = [d.dim_value if (d.HasField("dim_value")) else "?" for d in input.type.tensor_type.shape.dim]
    print(f"{input.name}: {dims}")

print("\n=== Outputs ===")
for output in model.graph.output:
    dims = [d.dim_value if (d.HasField("dim_value")) else "?" for d in output.type.tensor_type.shape.dim]
    print(f"{output.name}: {dims}")

実際、変換よりも shape の思い込みで詰まることの方が多いので、ここはかなり大事です。

ONNX Runtime でエンコーダ / デコーダをテストする

TensorRT の前に、まず ONNX Runtime でエンコーダ / デコーダの動作を確認しておくと切り分けしやすくなります。

import onnxruntime as ort
import numpy as np
import cv2
import matplotlib.pyplot as plt
import time


def load_and_preprocess_image(path):
    image = cv2.imread(path)
    if image is None:
        raise FileNotFoundError(f"画像が見つかりません: {path}")
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (1280, 720))
    image = image.astype(np.float32) / 255.0
    image = np.transpose(image, (2, 0, 1))
    return image


def run_encoder(providers, model_path, image_chw):
    sess_opt = ort.SessionOptions()
    sess_opt.log_severity_level = 0
    session = ort.InferenceSession(model_path, sess_options=sess_opt, providers=providers)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    start = time.perf_counter()
    output = session.run([output_name], {input_name: image_chw})
    end = time.perf_counter()
    print(f"[ENCODER] 推論時間: {(end - start) * 1000:.2f} ms")
    return output[0]


def run_decoder(providers, model_path, encoded):
    sess_opt = ort.SessionOptions()
    sess_opt.log_severity_level = 0
    session = ort.InferenceSession(model_path, sess_options=sess_opt, providers=providers)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    start = time.perf_counter()
    output = session.run([output_name], {input_name: encoded.astype(np.float32)})
    end = time.perf_counter()
    print(f"[DECODER] 推論時間: {(end - start) * 1000:.2f} ms")
    return output[0]


providers = [
    ('TensorrtExecutionProvider', {
        'device_id': 0,
        'trt_engine_cache_enable': True,
        'trt_engine_cache_path': './engine_cache',
        'trt_max_workspace_size': 2147483648,
        'trt_fp16_enable': True,
    }),
    ('CUDAExecutionProvider', {
        'device_id': 0,
        'arena_extend_strategy': 'kNextPowerOfTwo',
        'gpu_mem_limit': 4 * 1024 * 1024 * 1024,
        'cudnn_conv_algo_search': 'EXHAUSTIVE',
        'do_copy_in_default_stream': True,
    })
]

ここではコード全文を全部貼るより、「TensorRT EP と CUDA EP をどう並べて使っているか」がポイントです。

ONNX モデルを float16 化する

内部を軽量化したい場合は、float16 化も選択肢になります。

import onnx
from onnxconverter_common import float16

input_model = "model.onnx"
output_model = "model_fp16.onnx"

model = onnx.load(input_model)
fp16_model = float16.convert_float_to_float16(model, keep_io_types=True)
onnx.save(fp16_model, output_model)
print(f"Saved: {output_model}")

keep_io_types=True を付けると、入出力は float32 のまま、内部だけ float16 化 できます。これは後段の TensorRT でも扱いやすいことがあります。

確認には次を使います。

/usr/src/tensorrt/bin/trtexec --onnx=model_fp16.onnx

ONNX を SavedModel / TFLite 側へ戻す

逆方向の変換が必要な場合は onnx2tf が使えます。

pip install onnx onnxruntime
pip install tensorflow
pip install onnx2tf
pip install tf_keras onnx_graphsurgeon psutil ai_edge_litert sng4onnx

変換:

onnx2tf -i decoder13.onnx -o model_tf --output_signaturedefs

ノード名に問題がある場合は、記号を置換してから再保存することがあります。

import onnx

model = onnx.load("decoder.onnx")
for node in model.graph.node:
    node.name = node.name.replace("(", "_").replace(")", "_").replace("/", "_")
onnx.save(model, "decoder_clean.onnx")

INT8 量子化について

最終的に TFLite や TensorRT 側で軽量化したい場合、INT8 量子化も視野に入ります。

onnx2tf \
  -i encoder.onnx \
  -o ./output_dir_int8 \
  --quant_type int8 \
  --calib_data_dir ./calib_images

このあたりは画質・精度・速度のトレードオフが強いので、まずは float32 / float16 / INT8 のどこまで落とせるかを段階的に見るのが現実的です。

まとめ

Jetson で TensorRT を使う前に、まず整理したいのは次の点です。

• どのランタイム構成で行くか
• TFLite / TensorFlow 系モデルをどう ONNX にするか
• 変換後のノードや shape に問題がないか
• 必要に応じて onnxsim や float16 化を行うか

最初にこの整理をしておくと、後の engine 化や C++ 実装でかなり詰まりにくくなります。

次は、Jetson 上で ONNX Runtime を C++ から扱う準備を整理します。