Claude CodeでPyTorchモデルをONNXに書き出し、出力の一致をテストで検証する

PyTorchモデルをONNXに書き出す作業をClaude Codeに任せるなら、完了の条件は「.onnxファイルができた」ではなく「元のモデルと同じ入力で、onnxruntimeの出力が許容誤差に収まった」に置くのが扱いやすい構成です。ファイルは壊れていても生成されます。出力の一致をpytestで判定できれば、Claudeは書き出しコードを直すたびに同じコマンドで合否を読めます。この記事では、先に書かせるテスト、書き出しスクリプト、dynamic_shapesの落とし穴、CLAUDE.mdとStop hookの設定を順に示します。

完了の合図を「出力の一致」に置く理由

Claude Codeのベストプラクティスは、Claudeに合否が返ってくるチェックを渡すことを勧めています。チェックがないと、Claudeは「できたように見える」ところで止まり、間違いに気づくのは人間の役目になるためです。チェックは、テストファイルでも、出力を期待値と比べるスクリプトでもかまいません。

ONNXの書き出しは、このチェックが特に効く作業です。torch.onnx.exportは、形状や演算子の変換がうまくいけばファイルを返します。そのファイルが元のモデルと同じ値を出すかどうかは、別に確かめない限り分かりません。複数GPUの学習ジョブが止まる場合は、torchrunのNCCL timeoutを追う手順にまとめています。

確かめる対象は2つに分けられます。

• 同じ入力での出力差: 数値が許容誤差に収まっているか
• 入力形状を変えたときの挙動: 例に使ったバッチサイズ以外でも動くか

後者は、書き出し時にdynamic_shapesを渡し忘れると外れます。PyTorchのFAQにも、書き出したモデルの入力サイズが固定されているという質問が載っています。トレーサーは例に渡した入力の形状を記録するため、可変にしたい次元はdynamic_shapesで指定する必要があるという答えです。

前提と依存パッケージ

ONNXエクスポーターは、PyTorch本体に加えて追加のPythonパッケージを使います。PyTorchのドキュメントが挙げているのはonnxとonnxscriptで、書き出したモデルの実行にはonnxruntimeを使えるとされています。

pip install --upgrade torch onnx onnxscript onnxruntime

pip installはClaudeに実行させず、自分で済ませておく運用を想定しています。後述の権限設定では、pytestと書き出しスクリプトだけを許可し、パッケージの追加は許可しません。

torch.onnx.exportの既定は、PyTorch 2.9でdynamo=Trueになりました。torch.exportベースの新しい書き出し経路です。以下のコードはすべてこの経路を前提にします。古い記事やサンプルに出てくるdynamic_axesは、dynamo=Trueのときはdynamic_shapesを使うようPyTorchが案内しています。

先に書かせるテスト: PyTorchとonnxruntimeの出力を比べる

実装より先に、テストを書かせます。順序を逆にすると、Claudeは実装が出力した値に合わせて期待値を作りがちです。先にテストがあれば、書き出しスクリプトは「このテストを通す」ために書かれます。

比較の基準は、書き出す前のnn.Moduleにします。次のテストは、バッチサイズを変えながら、eval()にした元モデルの出力とonnxruntimeの出力を突き合わせます。

# tests/test_onnx_parity.py
import onnxruntime as ort
import pytest
import torch
 
from export_onnx import build_model, export
 
ATOL = 1e-4   # float32の出発点。緩めるときは理由を残す
RTOL = 1e-3
 
 
@pytest.fixture(scope="module")
def pair(tmp_path_factory):
    path = tmp_path_factory.mktemp("onnx") / "model.onnx"
    model = build_model().eval()
    export(model, path)
    session = ort.InferenceSession(
        str(path), providers=["CPUExecutionProvider"]
    )
    return model, session
 
 
@pytest.mark.parametrize("batch", [1, 3, 97, 256])
def test_output_matches_pytorch(pair, batch):
    model, session = pair
    x = torch.randn(batch, 8)
    with torch.no_grad():
        expected = model(x)
    (actual,) = session.run(None, {"x": x.numpy()})
    torch.testing.assert_close(
        torch.from_numpy(actual), expected, atol=ATOL, rtol=RTOL
    )

許容誤差の数値は、check_discrepanciesのドキュメントを目安にしています。atolの推奨値として、floatで1e-4、float16で1e-2が挙がっています。RTOLの1e-3は、ここで置いた出発点です。公式の値ではありません。

バッチサイズに1, 3, 97, 256を並べたのは意図があります。書き出し時の例にあたる値だけでなく、それより小さい値、大きい値、中間の値を流して、形状の固定漏れを拾うためです。

書き出しスクリプトとdynamic_shapes

テストが落ちることを確認してから、書き出しスクリプトを書かせます。PyTorchのMLP例に倣うと、次の形になります。

# export_onnx.py
import torch
import torch.nn as nn
 
 
def build_model() -> nn.Module:
    return nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 2))
 
 
def export(model: nn.Module, path) -> None:
    example = torch.randn(97, 8)
    program = torch.onnx.export(
        model,
        (example,),
        input_names=["x"],
        output_names=["y"],
        dynamic_shapes=({0: "batch"},),
        dynamo=True,
    )
    program.save(str(path))

ここに出てくる引数のうち、検証に関わるものを表にします。

dynamic_shapesの指定は、argsの入れ子の構造と対応させます。PyTorchの例では、位置引数、辞書、リストのそれぞれに{0: "batch_size"}を割り当てています。書き出し後にはonnx_program.model.graph.inputs[0].shapeが("batch_size", 8)になることをassertで確かめる例も載っており、これはテストにも流用できます。

PyTorch標準の検証APIとの使い分け

PyTorchには、書き出したモデルを検証するAPIが用意されています。自作のテストと役割が重なるため、違いを表にしておきます。

3つとも、比較の基準は元のnn.Moduleではなく、torch.exportが作ったExportedProgramです。書き出しの手前までに生じた差、たとえばtorch.exportの段階で固定された値は、この比較には現れません。自作テストが元のモデルと比べるのは、この部分を拾うためです。

使い分けは次のように考えられます。

• 自作のpytest: 完了条件。元のモデルとonnxruntimeを、複数のバッチサイズで比べる
• assert_onnx_program: 書き出し側のバグ切り分け。ExportedProgramとの差なのか、それ以前の差なのかを見分ける
• verify_onnx_program: 数値の確認。max_abs_diffを出力させ、許容誤差を決める根拠にする

テストが落ちたとき、Claudeに「verify_onnx_programの結果も出して原因を絞って」と指示すると、差が書き出し以降で生まれたのか、それ以前から存在したのかを分けて調べられます。

動的軸まわりの落とし穴

dynamic_shapesに関するつまずきは、症状から原因を引けるようにしておくと速く直せます。

変換そのものが失敗する場合は、torch.onnx.exportをreport=True付きでもう一度呼びます。診断用のMarkdownレポートがartifacts_dir(既定はカレントディレクトリ)に出るため、Claudeにそのファイルを読ませると、どの演算子で止まったかを手がかりにできます。

Claude Codeに任せる設定

ここまでの内容を、Claudeが毎回同じ手順で進められるように設定へ落とします。

CLAUDE.mdに書く内容

ベストプラクティスは、CLAUDE.mdを短く保つこと、特に検証コマンドを書くことを勧めています。コードから読み取れる情報は書かず、読み取れない決めごとだけを残します。

# Commands
- ONNX一致テスト: python -m pytest tests/test_onnx_parity.py -q
 
# Conventions
- 書き出しは torch.onnx.export(..., dynamo=True) と dynamic_shapes を使う。dynamic_axes は使わない
- 書き出す前に model.eval() を呼ぶ
- 許容誤差(ATOL/RTOL)を緩めてテストを通さない。緩める提案は理由を添えて人に確認する

最後の行は、運用上の保険です。出力が合わないとき、Claudeがテストを通すためにATOLを大きくする手もあります。この行は、その近道を人間の判断に回すためのものです。

Stop hookで合格まで終わらせない

CLAUDE.mdの記述は、守られないことがあります。pytest規約の書き方はCLAUDE.mdにpytestの規約を教える記事にあります。テストが通るまで終わらせたくない場合は、Stop hookが向いています。ベストプラクティスにあるtscの例を、このテストに置き換えると次のようになります。

{
  "hooks": {
    "Stop": [
      {
        "hooks": [
          {
            "type": "command",
            "command": "python -m pytest tests/test_onnx_parity.py -q 1>&2 || exit 2"
          }
        ]
      }
    ]
  }
}

終了コード2がClaudeの停止を止め、標準エラーに出した内容が「まだ終われない理由」としてClaudeに戻ります。pytestの出力を1>&2でそちらへ流しているのは、そのためです。

連続でブロックできる回数には上限があります。ブロックが8回続くと、Claude Codeが次のブロックを無視してターンを終わらせます。上限はCLAUDE_CODE_STOP_HOOK_BLOCK_CAPで変えられます。仕組みの詳細はStop hookのブロック上限の記事に、PostToolUseとの使い分けはhooksでテストを自動実行する記事にまとめています。

許可する操作を絞る

権限は、テストの実行と書き出しスクリプトだけに絞ります。ルールの書式はBash(npm run *)のように、サブコマンドの後ろに*を置く形です。

{
  "permissions": {
    "allow": [
      "Bash(python -m pytest *)",
      "Bash(python export_onnx.py *)"
    ]
  }
}

pip installは含めません。依存の追加は人が決めるものとして残します。

最初の指示文の例

セッションの冒頭では、順序と証拠の提出を指定します。ベストプラクティスも、Claudeに「実行したコマンドとその出力」を求めることを勧めています。

models/mlp.py のモデルをONNXに書き出す。
1. 先に tests/test_onnx_parity.py を書く。バッチ1・3・97・256で、元モデルとonnxruntimeの出力を assert_close で比べる
2. export_onnx.py がまだ無い状態でテストが落ちることを確認する
3. export_onnx.py を実装する。dynamic_shapes でバッチ次元を可変にする
4. 通ったら、実行したコマンドと出力をそのまま貼る

pytestの2回目以降は、Claudeが出力を読んで直す反復になります。人が見るのは、最後に貼られた出力と、許容誤差が変わっていないかだけです。

テスト範囲を広げるときの目安

最初のテストは、バッチ次元だけを可変にして単純な入力を流します。モデルの実態に近づけるときは、次の順で足していくと原因を切り分けやすくなります。

1. 可変にしたい次元を増やす(系列長、画像の高さ・幅など)。dynamic_shapesに次元を足し、parametrizeに値を足す
2. 実データに近い入力を流す。乱数だけでは通る分岐が、実データでは別の経路を通ることがあります
3. 許容誤差を精度に合わせる。float16を使うなら1e-2が目安です
4. 出力が複数あるモデルは、出力ごとにassert_closeを呼ぶ

制御フローを含むモデルは、書き出し自体が難しくなります。FAQは、ループを含むモデルの書き出しについてtorch.condを参照するよう案内しています。学習や推論そのものが遅いときの切り分けは、torch.profilerでDataLoaderの待ちを直す手順にあります。

まとめ

ONNXの書き出しでは、ファイルができることと、値が元のモデルと一致することは別の問題です。完了条件を後者に置き、元のnn.Moduleとonnxruntimeを複数のバッチサイズで比べるテストを先に書かせると、Claudeは書き出しコードの修正を自分で回せます。

最初に確かめたいのは、例に使わなかったバッチサイズでテストが通るかどうかです。そこが落ちるなら、dynamic_shapesの指定漏れを疑う順序で進められます。