エラー解決

【PyTorch】ONNXエクスポートとTensorRT変換の完全ガイド：よくあるエラーと解決策

· 18分で読めます · 2 閲覧

問題の概要：PyTorchモデルのONNXエクスポートとTensorRT変換における課題

PyTorchで学習したモデルを本番環境で高速に推論するためには、ONNX形式へのエクスポートと、それをNVIDIAのTensorRTエンジンに変換するプロセスが一般的です。しかし、この一連の作業では、初心者から中級者までが直面する特有のエラーや課題が数多く存在します。具体的には、ONNXエクスポート時の「シンボリックサイズ」に関するエラー、TensorRT変換時のプラグイン不足エラー、そして変換後のモデルで精度が低下する問題などが代表的です。本記事では、これらの課題を具体的なエラーメッセージと共に解説し、実践的な解決手順を提供します。

原因の解説

ONNX（Open Neural Network Exchange）は、異なるディープラーニングフレームワーク間でモデルを交換するためのオープンなフォーマットです。TensorRTはNVIDIA GPU上で推論を最適化・高速化する推論エンジンです。変換プロセスが複雑な理由は以下の通りです。

1. 動的シェープと静的シェープの不一致

PyTorchモデルは、バッチサイズやシーケンス長などが動的（可変）であることが多いです。しかし、TensorRTは多くの場合、最適化のために入出力テンソルの形状（シェープ）を静的に固定することを要求します。ONNXエクスポート時に動的次元を適切に指定しないと、後続の変換で失敗します。

2. 演算（OP）のサポート範囲の違い

PyTorch、ONNX、TensorRTはそれぞれサポートする演算の種類やバージョンが異なります。PyTorchにしかないカスタム演算や、ONNX/TensorRTがまだサポートしていない新しい演算を使用していると、変換に失敗します。

3. バージョン間の非互換性

PyTorch、torch.onnx、onnx、TensorRT（trtexec）の各ライブラリ・ツールのバージョンの組み合わせによって、動作が大きく変わることがあります。公式にテストされていない組み合わせでは予期せぬエラーが発生します。

解決方法：ステップバイステップ手順

ステップ1: 環境構築とバージョン確認

まず、互換性のあるバージョンを揃えることが最も重要です。以下は、2023年後半時点で比較的安定した組み合わせの一例です。

# 推奨バージョン例
PyTorch: 1.13.1 または 2.0.1
torchvision: 対応するバージョン
ONNX: 1.14.0
onnxruntime: 1.15.0
TensorRT: 8.6.1
CUDA: 11.8
cuDNN: 8.9

ステップ2: PyTorchモデルのONNXエクスポート（動的シェープ対応）

最も一般的なエラー「Exporting the operator … to ONNX opset version … is not supported.」や「Symbolic size … must be …」を避けるために、動的次元を明示的に指定します。

import torch
import torch.onnx

# サンプルモデル（例: 簡単なCNN）
class SimpleCNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.fc = torch.nn.Linear(16*32*32, 10)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleCNN()
model.eval()  # 推論モードに必須

# ダミー入力
dummy_input = torch.randn(1, 3, 32, 32)

# 動的シェープを指定してONNXエクスポート
# ここではバッチサイズを動的(-1)に、その他は固定とします
dynamic_axes = {
    'input': {0: 'batch_size'},   # 入力の0次元（バッチ）を動的に
    'output': {0: 'batch_size'}   # 出力も同様
}

onnx_path = "model.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes,
    opset_version=14,  # 安定したopsetを指定（11, 13, 14が一般的）
    do_constant_folding=True,  # 定数畳み込み最適化を有効化
    verbose=False
)

print(f"モデルを {onnx_path} にエクスポートしました。")

ステップ3: ONNXモデルの検証と簡略化

エクスポートしたONNXモデルに問題がないか検証し、必要に応じて簡略化します。

import onnx
import onnxruntime as ort

# ONNXモデルのロードと検証
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNXモデルの構文チェックが成功しました。")

# ONNX Runtimeで推論テスト（変換の検証）
ort_session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
outputs = ort_session.run(None, {'input': dummy_input.numpy()})
print("ONNX Runtimeでの推論テストが成功しました。")

# モデルの簡略化（オプションだが推奨）
# onnx-simplifier をインストール必要: pip install onnx-simplifier
# import onnxsim
# model_simp, check = onnxsim.simplify(onnx_model)
# onnx.save(model_simp, "model_simplified.onnx")

ステップ4: TensorRTエンジンへの変換（trtexec使用）

ONNXモデルをTensorRTエンジン（.planまたは.engineファイル）に変換します。ここで「Could not find plugin for node: …」といったエラーが発生することがあります。

方法A: trtexecコマンドラインツールを使用（推奨）

# 基本的な変換コマンド
# TensorRTのインストールパスに`trtexec`があります
trtexec --onnx=model.onnx 
        --saveEngine=model.engine 
        --workspace=2048 
        --fp16  # FP16精度を有効化（Volta以降のGPUで推奨）

# 動的シェープを指定する場合（バッチサイズが動的）
trtexec --onnx=model.onnx 
        --saveEngine=model_dynamic.engine 
        --minShapes=input:1x3x32x32 
        --optShapes=input:8x3x32x32 
        --maxShapes=input:32x3x32x32 
        --workspace=2048 
        --fp16

# ビルドログを詳細に出力してエラーを確認
trtexec --onnx=model.onnx --verbose 2>&1 | tee build.log

方法B: Python APIを使用（より細かい制御が必要な場合）

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))  # エラー詳細を表示

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2048 * 1024 * 1024)  # 2GB
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

# 動的シェーププロファイルの設定（動的バッチの場合）
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 32, 32), (8, 3, 32, 32), (32, 3, 32, 32))
config.add_optimization_profile(profile)

serialized_engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(serialized_engine)
print("TensorRTエンジンのビルドが完了しました。")

ステップ5: 変換後の精度検証

変換前後のモデルで推論結果を比較し、精度の大幅な低下がないことを確認します。

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

# TensorRTエンジンのロードと実行
def load_engine(engine_path):
    with open(engine_path, "rb") as f:
        runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
        engine = runtime.deserialize_cuda_engine(f.read())
    return engine

engine = load_engine("model.engine")
context = engine.create_execution_context()

# 入出力バッファの設定
dummy_input_np = np.random.randn(1, 3, 32, 32).astype(np.float32)

# メモリ割り当てと転送（簡略版。実際は非同期処理やストリーム管理が必要）
output = np.empty((1, 10), dtype=np.float32)
d_input = cuda.mem_alloc(dummy_input_np.nbytes)
d_output = cuda.mem_alloc(output.nbytes)

cuda.memcpy_htod(d_input, dummy_input_np)
context.execute_v2([int(d_input), int(d_output)])
cuda.memcpy_dtoh(output, d_output)

print("TensorRT推論結果の形状:", output.shape)

# PyTorchの元の結果と比較（簡易チェック）
with torch.no_grad():
    torch_output = model(torch.from_numpy(dummy_input_np)).numpy()

diff = np.abs(output - torch_output).max()
print(f"PyTorchとTensorRTの出力最大差分: {diff}")
if diff < 1e-3:
    print("✅ 精度は許容範囲内です。")
else:
    print("⚠️  精度に大きな差があります。FP16モードをオフにするか、動的範囲の較正を検討してください。")

よくあるエラーとその解決策

エラー1: “Unsupported: ONNX export of … on a tensor of dtype …”

原因: ONNX opsetが特定のデータ型の演算をサポートしていない。

解決策: opset_versionを変更する（例: 11→13）。または、モデル内の該当演算をサポートされているものに置き換える。

エラー2: “TensorRT could not find plugin for node: …”

原因: ONNXモデルにTensorRTがネイティブサポートしていない演算が含まれている。

解決策: 以下のいずれかを行います。
1. ONNXモデルを簡略化（onnx-simplifier）して、サポートされない演算を除去する。
2. 該当するカスタムプラグインを実装・登録する。
3. モデルアーキテクチャを変更し、サポートされている演算のみを使用する。

エラー3: “Invalid ONNX opset: …”

原因: 使用しているTensorRTのバージョンが、ONNXモデルのopsetバージョンをサポートしていない。

解決策: torch.onnx.export() の opset_version を下げる（例: 14→13または11）。TensorRT 8.xではopset 13までが安定しています。

まとめ・補足情報

PyTorchモデルからONNXを経由してTensorRTエンジンへの変換は、本番環境での高速推論にはほぼ必須のプロセスです。成功の鍵は以下の3点に集約されます。

  1. バージョン管理: PyTorch、ONNX、TensorRTの互換性を慎重に確認し、可能な限り公式ドキュメントで推奨される組み合わせを使用する。
  2. 動的シェープの明示: エクスポート時に`dynamic_axes`を適切に設定し、TensorRT変換時にも対応するシェーププロファイルを設定する。
  3. 段階的な検証: PyTorch → ONNX → ONNX Runtime → TensorRT の各段階で推論結果を検証し、問題を早期に発見・切り分ける。

また、最新の機能として、PyTorch 2.0以降では`torch.compile`や`torch.export`といった新しいエクスポートパスも登場しています。しかし、現時点ではエコシステムの成熟度と安定性を考慮すると、ここで解説した従来の`torch.onnx.export`を用いたワークフローが最も信頼性が高いと言えます。変換に失敗した場合は、エラーメッセージを仔細に読み、モデルを可能な限りシンプルにすることから始めると良いでしょう。

この記事は役に立ちましたか？