PyTorchでONNXモデルを書き出す際、大きなモデルではONNX External Data形式を利用して、モデル本体と重みファイルを別々に保存したいことがあります。

この記事では、

• model.onnx
• model_weights.bin

の2ファイルに分けて出力する方法と、古いPyTorchを使用している場合の対処法を紹介します。

通常のONNX出力

一般的なONNX出力は次のようになります。

torch.onnx.export(
    wrapper,
    (sample, timestep, encoder_hidden_states),
    os.path.join(outdir, "unet", "model.onnx"),

    input_names=["sample", "timestep", "encoder_hidden_states"],
    output_names=["out_sample"],

    dynamic_axes={
        "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
        "timestep": {0: "batch"},
        "encoder_hidden_states": {0: "batch", 1: "sequence"},
        "out_sample": {
            0: "Convout_sample_dim_0",
            2: "Convout_sample_dim_2",
            3: "Convout_sample_dim_3",
        },
    },

    do_constant_folding=True,
    opset_version=opset,
)

このままでは通常、1つの model.onnx が出力されます。

方法1：PyTorch 2.xならExternal Data形式で出力する

PyTorchの新しいバージョンでは、use_external_data_formatを指定するだけで重みを別ファイルにできます。

torch.onnx.export(
    wrapper,
    (sample, timestep, encoder_hidden_states),
    os.path.join(outdir, "unet", "model.onnx"),

    input_names=["sample", "timestep", "encoder_hidden_states"],
    output_names=["out_sample"],

    dynamic_axes={
        "sample": {0: "batch", 1: "channels", 2: "height", 3: "width"},
        "timestep": {0: "batch"},
        "encoder_hidden_states": {0: "batch", 1: "sequence"},
        "out_sample": {
            0: "Convout_sample_dim_0",
            2: "Convout_sample_dim_2",
            3: "Convout_sample_dim_3",
        },
    },

    do_constant_folding=True,
    opset_version=opset,

    use_external_data_format=True,
)

出力例

unet/
├─ model.onnx
└─ model.onnx.data

環境によっては

unet/
├─ model.onnx
└─ model_weights.bin

となります。

重みファイル名を指定する

PyTorch 2.xでは、external_data_pathを指定できます。

torch.onnx.export(
    wrapper,
    (sample, timestep, encoder_hidden_states),
    os.path.join(outdir, "unet", "model.onnx"),

    input_names=["sample", "timestep", "encoder_hidden_states"],
    output_names=["out_sample"],

    dynamic_axes={...},

    do_constant_folding=True,
    opset_version=opset,

    use_external_data_format=True,
    external_data_path="model_weights.bin",
)

これで

unet/
├─ model.onnx
└─ model_weights.bin

という構成になります。

TypeErrorが出る場合

もし次のようなエラーが表示された場合

TypeError: export() got an unexpected keyword argument 'use_external_data_format'

使用しているPyTorchが古い可能性があります。

PyTorchのバージョンを確認する

import torch

print(torch.__version__)

目安は以下のとおりです。

可能であれば最新版へアップデートするのがおすすめです。

pip install -U torch

方法2：古いPyTorchなら後からExternal Dataへ変換する

環境の都合でPyTorchを更新できない場合でも、ONNXライブラリを使えば対応できます。

① 通常通りONNXを書き出す

import torch.onnx
import onnx

onnx_path = os.path.join(outdir, "unet", "model.onnx")

torch.onnx.export(
    wrapper,
    (sample, timestep, encoder_hidden_states),
    onnx_path,
    input_names=["sample", "timestep", "encoder_hidden_states"],
    output_names=["out_sample"],
    dynamic_axes={...},
    opset_version=opset,
)

② External Data形式へ保存し直す

import onnx

model = onnx.load(onnx_path)

onnx.save_model(
    model,
    onnx_path,
    save_as_external_data=True,
    all_tensors_to_one_file=True,
    location="model_weights.bin",
)

出力結果

unet/
├─ model.onnx
└─ model_weights.bin

この方法ならPyTorchのバージョンに依存せず利用できます。

ONNX Runtimeでの読み込み

外部データ形式でも読み込み方法は通常と同じです。

import onnxruntime as ort

sess = ort.InferenceSession("model.onnx")

model_weights.binは自動的に読み込まれます。

不要なファイルを削除する方法

変換後にmodel.onnxとmodel_weights.binだけ残したい場合は、以下のコードが便利です。

import os

target_dir = "/path/to/your/folder"

keep_files = {"model.onnx", "model_weights.bin"}

for name in os.listdir(target_dir):
    path = os.path.join(target_dir, name)

    if name in keep_files:
        continue

    if os.path.isfile(path):
        os.remove(path)

サブフォルダも削除する場合

import os
import shutil

target_dir = "/path/to/your/folder"

keep_files = {"model.onnx", "model_weights.bin"}

for name in os.listdir(target_dir):
    path = os.path.join(target_dir, name)

    if name in keep_files:
        continue

    if os.path.isfile(path):
        os.remove(path)
    elif os.path.isdir(path):
        shutil.rmtree(path)

注意: このコードはサブフォルダも完全に削除します。

まとめ

今回紹介した内容をまとめると、次のようになります。

• PyTorch 2.xならuse_external_data_format=Trueを利用する
• external_data_pathで重みファイル名を指定できる
• 古いPyTorchではonnx.save_model()でExternal Data形式へ変換できる
• ONNX RuntimeやTensorRTでもそのまま利用可能
• 変換後は不要なファイルをPythonで自動削除できる

PyTorchのバージョンによって対応方法が異なるため、自分の環境に合った方法を選択しましょう。