1
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@GeneLab_999

【第3回】実践！PyTorchモデルをONNXに変換してみよう

1
Posted at

ねらい

PyTorchで作成したモデルをONNX形式に変換する方法を実際に手を動かして理解する。

対象

  • PyTorchでモデルを訓練したことがある人
  • モデルをデプロイしたい人
  • 「ONNXに変換したいけど、やり方がわからない」人

ゴール

  • torch.onnx.export() を使いこなせるようになる
  • 動的入力サイズへの対応ができる
  • 変換後のモデルを検証できる

TL;DR

  • torch.onnx.export() で変換する
  • ダミー入力が必要（トレーシングのため）
  • dynamic_axes で可変サイズに対応
  • onnx.checker.check_model() で検証

環境準備

まずは必要なライブラリをインストールしよう。

pip install torch onnx onnxruntime

バージョン確認：

import torch
import onnx
import onnxruntime as ort

print(f"PyTorch: {torch.__version__}")
print(f"ONNX: {onnx.__version__}")
print(f"ONNX Runtime: {ort.__version__}")

執筆時点での推奨バージョン：

  • PyTorch 2.0以上
  • ONNX 1.14以上
  • ONNX Runtime 1.15以上

最もシンプルな変換

まずは基本中の基本から。シンプルなMLPを変換してみよう。

モデルの定義

import torch
import torch.nn as nn

class SimpleMLP(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=256, output_dim=10):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# モデルのインスタンス化
model = SimpleMLP()
model.eval()  # 推論モードに設定（重要！）

ONNX変換

# ダミー入力を作成（トレーシングに必要）
dummy_input = torch.randn(1, 784)

# ONNXにエクスポート
torch.onnx.export(
    model,                    # 変換するモデル
    dummy_input,              # ダミー入力
    "simple_mlp.onnx",        # 出力ファイル名
    export_params=True,       # 重みも一緒に保存
    opset_version=11,         # opset version（11以上推奨）
    input_names=["input"],    # 入力の名前
    output_names=["output"]   # 出力の名前
)

print("変換完了！")

たったこれだけ。simple_mlp.onnx というファイルが生成される。

変換結果の検証

import onnx

# モデルを読み込み
onnx_model = onnx.load("simple_mlp.onnx")

# モデルの検証（不正な場合は例外が発生）
onnx.checker.check_model(onnx_model)
print("モデルは正常です")

# グラフの情報を表示
print(f"グラフ名: {onnx_model.graph.name}")
print(f"ノード数: {len(onnx_model.graph.node)}")

torch.onnx.export() のパラメータ詳解

公式ドキュメントから重要な部分を見てみよう。

Exporting a model in PyTorch works via tracing or scripting. This tutorial will use as an example a model exported by tracing. To export a model, we call the torch.onnx.export() function. This will execute the model, recording a trace of what operators are used to compute the outputs.

（PyTorchでのモデルのエクスポートは、トレーシングまたはスクリプティングを通じて動作する。このチュートリアルでは、トレーシングによってエクスポートされたモデルを例として使用する。モデルをエクスポートするには、torch.onnx.export()関数を呼び出す。これにより、モデルが実行され、出力を計算するために使用されるオペレータのトレースが記録される。）

出典: PyTorch ONNX Tutorial

主要なパラメータ

torch.onnx.export(
    model,                    # nn.Module または ScriptModule
    args,                     # ダミー入力（タプルまたはテンソル）
    f,                        # 出力ファイルパス
    export_params=True,       # 学習済みパラメータを含めるか
    verbose=False,            # 詳細なログを出力するか
    training=TrainingMode.EVAL,  # 訓練モード or 推論モード
    input_names=None,         # 入力テンソルの名前
    output_names=None,        # 出力テンソルの名前
    opset_version=None,       # ONNXのopset version
    dynamic_axes=None,        # 動的な次元の指定
    do_constant_folding=True, # 定数畳み込み最適化
)

特に重要なパラメータ

export_params

True にすると学習済みの重みがモデルに含まれる。本番用なら True 一択。

opset_version

前回説明したopset version。明示的に指定しよう。

# 推奨: 11以上を明示的に指定
opset_version=11

do_constant_folding

定数の事前計算を行う最適化。通常は True でOK。

# 例: a = 2 * 3 のような定数計算を、変換時に a = 6 に置き換える
do_constant_folding=True

動的入力サイズへの対応

実運用では、バッチサイズや入力サイズが変わることが多い。固定サイズだと不便だ。

Note that the input size will be fixed in the exported ONNX graph for all the input's dimensions, unless specified as a dynamic axes.

（入力サイズは、動的軸として指定されない限り、エクスポートされたONNXグラフ内のすべての入力次元で固定されることに注意。）

出典: PyTorch ONNX Tutorial

dynamic_axes の使い方

# バッチサイズを可変にする
torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    opset_version=11,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},   # 0番目の次元を可変に
        "output": {0: "batch_size"}   # 出力も同様
    }
)

複数の可変次元

画像処理モデルなら、高さと幅も可変にしたいことがある。

# 画像モデル用（バッチ、高さ、幅を可変に）
# 入力形状: [batch, channels, height, width]
dynamic_axes={
    "input": {
        0: "batch_size",
        2: "height",
        3: "width"
    },
    "output": {
        0: "batch_size"
    }
}

シーケンス長を可変にする（NLP）

# テキストモデル用
# 入力形状: [batch, sequence_length]
dynamic_axes={
    "input_ids": {
        0: "batch_size",
        1: "sequence_length"
    },
    "attention_mask": {
        0: "batch_size",
        1: "sequence_length"
    },
    "output": {
        0: "batch_size"
    }
}

画像分類モデルの変換例

実践的な例として、ResNet18を変換してみよう。

import torch
import torchvision.models as models

# 学習済みResNet18を読み込み
model = models.resnet18(pretrained=True)
model.eval()

# ダミー入力（バッチサイズ1、3チャンネル、224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# ONNXにエクスポート
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=["image"],
    output_names=["class_logits"],
    dynamic_axes={
        "image": {0: "batch_size"},
        "class_logits": {0: "batch_size"}
    }
)

print("ResNet18をONNXに変換しました")

変換結果の確認

import onnx
from onnx import shape_inference

# モデルを読み込んで検証
model = onnx.load("resnet18.onnx")
onnx.checker.check_model(model)

# 形状推論を実行（デバッグに便利）
model_with_shapes = shape_inference.infer_shapes(model)

# 入出力の情報を表示
for input in model.graph.input:
    print(f"入力: {input.name}")
    
for output in model.graph.output:
    print(f"出力: {output.name}")

ONNX Runtimeで推論してみる

変換したモデルが正しく動くか確認しよう。

import numpy as np
import onnxruntime as ort

# ONNX Runtimeセッションを作成
session = ort.InferenceSession("resnet18.onnx")

# 入力情報を確認
input_name = session.get_inputs()[0].name
print(f"入力名: {input_name}")

# ダミーデータで推論
dummy_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {input_name: dummy_data})

print(f"出力形状: {outputs[0].shape}")
print(f"出力サンプル: {outputs[0][0, :5]}")  # 最初の5クラスの値

PyTorchの出力と比較

変換が正しく行われたか、数値的に確認する。

import torch
import numpy as np
import onnxruntime as ort
import torchvision.models as models

# PyTorchモデル
torch_model = models.resnet18(pretrained=True)
torch_model.eval()

# ONNXモデル
ort_session = ort.InferenceSession("resnet18.onnx")

# 同じ入力データを作成
np_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
torch_input = torch.from_numpy(np_input)

# PyTorchで推論
with torch.no_grad():
    torch_output = torch_model(torch_input).numpy()

# ONNXで推論
ort_output = ort_session.run(None, {"image": np_input})[0]

# 差分を確認
diff = np.abs(torch_output - ort_output)
print(f"最大差: {diff.max()}")
print(f"平均差: {diff.mean()}")

# 許容誤差内か確認
if np.allclose(torch_output, ort_output, rtol=1e-5, atol=1e-6):
    print("変換成功！出力が一致しています")
else:
    print("警告: 出力に差異があります")

通常、差異は非常に小さい（1e-6程度）はずだ。

よくあるエラーと対処法

エラー1: TracingWarning

TracingWarning: Converting a tensor to a Python boolean might cause 
the trace to be incorrect.

原因: トレース時に、テンソルの値に依存する分岐がある。

# NG: 値に依存する分岐
def forward(self, x):
    if x.sum() > 0:  # テンソルの値で分岐
        return x * 2
    else:
        return x

対処:

  • 条件分岐を避ける
  • torch.where() を使う
  • torch.jit.script を使う
# OK: torch.whereを使う
def forward(self, x):
    return torch.where(x.sum() > 0, x * 2, x)

エラー2: Unsupported operator

RuntimeError: Exporting the operator 'aten::xxx' to ONNX opset version 11 
is not supported.

原因: そのオペレータがONNXでサポートされていない。

対処:

  • opset versionを上げてみる（サポートされていることがある）
  • サポートされている演算で書き換える
  • カスタムオペレータを登録する（上級者向け）

エラー3: Dynamic control flow

RuntimeError: Unsupported: ONNX export of operator ... with dynamic inputs

原因: 入力サイズに依存するループや分岐がある。

対処:

  • torch.onnx.export の代わりに torch.onnx.dynamo_export を試す（PyTorch 2.0以降）
  • モデル構造を見直す

エラー4: eval()を忘れた

症状: 推論結果がおかしい、バッチノーマライゼーションが効いていない。

原因: model.eval() を呼び忘れている。

対処: エクスポート前に必ず model.eval() を呼ぶ。

model.eval()  # これ忘れがち！
torch.onnx.export(model, ...)

実践Tips

1. verboseオプションを活用

問題が起きたら、まず詳細ログを見る。

torch.onnx.export(
    model, dummy_input, "model.onnx",
    verbose=True  # 詳細なログを出力
)

2. 段階的に確認

複雑なモデルは、サブモジュールごとに変換を試す。

# エンコーダだけ先に変換してみる
torch.onnx.export(model.encoder, dummy_input, "encoder.onnx", ...)

3. Netronで可視化

変換後は必ずNetronで構造を確認する習慣をつけよう。

# Netron（https://netron.app/）にドラッグ&ドロップ

4. 入出力の名前は意味のあるものに

デバッグ時に役立つ。

# NG
input_names=["input_0"]

# OK
input_names=["image_rgb"]

まとめ

  • torch.onnx.export() で変換する
  • ダミー入力を忘れずに用意
  • model.eval() を呼んでから変換
  • opset_version は明示的に指定（11以上推奨）
  • dynamic_axes で可変サイズに対応
  • 変換後は onnx.checker.check_model() で検証
  • PyTorchとONNXの出力を比較して確認

次回予告

第4回では、ONNX Runtimeを使った高速推論の実装を解説する。CPU/GPUの使い分け、IOBindingによる最適化、実際のベンチマーク結果も紹介する予定だ。

参考文献

1
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
1
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?