NVIDIA TensorRT Model Optimizerとは
TensorRT Model Optimizerとは、量子化、蒸留、プルーニング、推測的デコード、スパース性など、モデルの圧縮技術を利用して、モデルを最適化するPythonライブラリです。
Hugging Face、PyTorch、ONNXのモデルに対応しています。
Model OptimizerでONNXを混合精度に変換する
モデルの最適化は、LLMや画像生成など、大規模なモデルでは有効ですが、ここではONNXを混合精度にするためにModel OptimizerのAutoCastという機能を利用します。
ONNXの入出力はFP32のままにしたいがONNX内部ではFP16にしたい、FP32なコードをFP16に直すのがめんどくさい、というわたし向きな機能です。
Model Optimizer AutoCast の使い方
1.セットアップ
Model Optimizerは、pipでセットアップできるのですが、利用したいAutoCastはpipではインストールされません。
githubからクローンして利用します。
# Clone the Model Optimizer repository
git clone git@github.com:NVIDIA/TensorRT-Model-Optimizer.git
cd TensorRT-Model-Optimizer
pip install -e .[dev]
2.使い方
基本的には公式ドキュメントやサンプルコードのコメントに記載されている通りです。
modelopt_autocast.py
import onnx
from modelopt.onnx.autocast import convert_to_mixed_precision
# Convert model to mixed precision
converted_model = convert_to_mixed_precision(
onnx_path="model.onnx",
low_precision_type="fp16", # or "bf16"
nodes_to_exclude=None, # optional list of node name patterns to keep in FP32
op_types_to_exclude=None, # optional list of op types to keep in FP32
data_max=512, # threshold for node outputs
init_max=65504, # threshold for initializers
keep_io_types=False, # whether to preserve input/output types
calibration_data=None, # optional path to input data file
init_conversion_max_bytes=None, # maximum size in bytes for initializer conversion
providers=["cpu"], # list of Execution Providers for ONNX-Runtime backend
trt_plugins=[], # list of TensorRT plugin library paths in .so format
max_depth_of_reduction=None, # maximum depth of reduction allowed in low precision
)
# Save the converted model
onnx.save(converted_model, "converted_model.onnx")
3.私の場合
私はdynamic=trueなyolov8のモデルを混合精度にしました。
onnx2fp16.py
import onnx
import numpy as np
from modelopt.onnx.autocast import convert_to_mixed_precision
# Create random calibration data with fixed shape and normalize to [0, 1]
calibration_data = {"images": np.random.rand(1, 3, 736, 1280).astype(np.float32)}
# Convert model to mixed precision
converted_model = convert_to_mixed_precision(
onnx_path="my_yolov8m_s_dy.onnx",
low_precision_type="fp16", # or "bf16"
nodes_to_exclude=None, # optional list of node name patterns to keep in FP32
op_types_to_exclude=["Range"], # optional list of op types to keep in FP32
data_max=512, # threshold for node outputs
init_max=65504, # threshold for initializers
keep_io_types=True, # whether to preserve input/output types
calibration_data=calibration_data, # random calibration data
init_conversion_max_bytes=None, # maximum size in bytes for initializer conversion
providers=["cuda:0"], # list of Execution Providers for ONNX-Runtime backend
trt_plugins=[], # list of TensorRT plugin library paths in .so format
max_depth_of_reduction=None, # maximum depth of reduction allowed in low precision
)
# Save the converted model
onnx.save(converted_model, "converted_model.onnx")
変更した個所だけ解説
- op_types_to_exclude は、NoneではRangeにおいてエラーが出たので、RangeはFP32としました
onnx.onnx_cpp2py_export.shape_inference.InferenceError: [ShapeInferenceError] (op_type:Range, node name: /model.22/Range): start typestr: T, has unsupported type: tensor(float16) - calibration_data として、ランダムデータを入力としています。dynamicなため入力が決まらないと後ろが決まらないためです
- keep_io_types は入出力をFP32にするためにTrueとしています
- providersはCPUのままでも良さそうですが、cuda:0を指定しました
- 私の場合、92%をFP16にできたようです
2025-09-28 11:05:32,763 - autocast - INFO - convert.py - Converted 359/389 nodes (92.29%) to fp16
まとめ
- ONNXをできるだけFP16にしたいけど、コードはFP32のままなので入出力だけFP32にしたいよ！というようなニーズに、TensorRT Model Optimizer AutoCastは有効です
- 同じような機能のものに、onnxconverter_common の auto_mixed_precision があります
https://onnxruntime.ai/docs/performance/model-optimizations/float16.html
私はTensorRT Model Optimizerを利用しましたが、お好みで使いやすい方を使っていただければと思います。