PyTorchのtorch.onnx.export()関数：詳細解説

PyTorch の ONNX と torch.onnx.export()

torch.onnx.export() は、以下の引数を受け取ります。

• model: 変換したい PyTorch モデル
• file_path: 出力する ONNX ファイルのパス
• input_names: モデルの入力の名前のリスト
• dynamic_axes: 動的軸の名前のリスト (オプション)
• export_params: モデルのパラメータを含めるかどうか (デフォルト: True)
• opset_version: 使用する ONNX オペレータセットのバージョン (デフォルト: 13)

例：

import torch

model = torch.nn.Sequential(
    torch.nn.Linear(10, 100),
    torch.nn.ReLU(),
    torch.nn.Linear(100, 10)
)

torch.onnx.export(
    model,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
)

このコードは、model を model.onnx という名前の ONNX ファイルに変換します。モデルの入力は input という名前で、出力は output という名前になります。

torch.onnx.export() には、以下のオプションがあります。

• input_shape: 入力データの形状 (オプション)

input_shape オプションは、モデルの入力データの形状を指定するために使用されます。動的軸を使用する場合は、dynamic_axes オプションでその名前を指定する必要があります。

export_params オプションは、モデルのパラメータを ONNX ファイルに含めるかどうかを指定します。デフォルトでは True になっていますが、推論時にのみモデルを使用する場合は False に設定することができます。

opset_version オプションは、使用する ONNX オペレータセットのバージョンを指定します。使用可能なバージョンは、ONNX 公式ドキュメント: https://onnx.ai/ で確認することができます。

• torch.onnx.export() は、すべての PyTorch モデルを ONNX 形式に変換できるわけではありません。
• 変換された ONNX モデルは、元の PyTorch モデルと同じ精度で動作するとは限りません。

まとめ

torch.onnx.export() は、PyTorch モデルを ONNX 形式に変換するための便利な関数です。この関数を使って、PyTorch モデルを他のフレームワークで利用したり、推論速度を向上させたりすることができます。

• 上記の解説は、基本的な内容のみを記載しています。詳細については、上記の参考資料を参照してください。
• 問題が発生した場合は、PyTorch の公式ドキュメントやフォーラムなどで解決策を探してください。

PyTorch ONNX export サンプルコード

画像分類モデル

import torch
from torchvision import models

model = models.resnet18(pretrained=True)

# ダミーの入力データを作成
input_data = torch.randn(1, 3, 224, 224)

# モデルを ONNX 形式に変換
torch.onnx.export(
    model,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
)

自然言語処理モデル

import torch
from transformers import BertModel

model = BertModel.from_pretrained("bert-base-uncased")

# ダミーの入力データを作成
input_ids = torch.randint(0, 30000, (1, 128))
attention_masks = torch.ones((1, 128))

# モデルを ONNX 形式に変換
torch.onnx.export(
    model,
    "bert-base-uncased.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["last_hidden_state"],
    dynamic_axes={"input_ids": {0: "batch_size"}, "attention_mask": {0: "batch_size"}},
)

生成モデル

import torch
from torch.distributions import Normal

class Generator(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mu = torch.nn.Parameter(torch.zeros(10))
        self.sigma = torch.nn.Parameter(torch.ones(10))

    def forward(self):
        return Normal(self.mu, self.sigma).rsample()

model = Generator()

# ダミーの入力データを作成
batch_size = 10

# モデルを ONNX 形式に変換
torch.onnx.export(
    model,
    "generator.onnx",
    input_names=[],
    output_names=["output"],
    dynamic_axes={"output": {0: "batch_size"}},
)

• 上記のサンプルコードは、PyTorch 1.9.1 で動作確認しています。

PyTorch モデルを ONNX 形式に変換する他の方法

ONNX Runtime には、onnxruntime.convert.py というスクリプトが付属しています。このスクリプトを使って、PyTorch モデルを ONNX 形式に変換することができます。

例：

python onnxruntime/convert/python/onnx_converter.py \
    --model_path model.pt \
    --output_path model.onnx \
    --input_names input \
    --output_names output

torch.onnx.export() 以外にも、PyTorch モデルを ONNX 形式に変換する方法はいくつかあります。それぞれの方法にはメリットとデメリットがあるので、目的に合った方法を選択する必要があります。