ステータス:ドラフト(2026-10-16)、一部実測。 torch.compile・メモリレイアウト・cuDNNの設定は、ジブラルタルの8GBクラスのノートPC用GPU1枚で使い捨てPodを使って測定しました(PyTorch 2.9、ウォームアップ20回の後、50回×5セットの中央値)。融合の原理はノートPCのCPUの例で確認しました(
code/03_fusion.py、数秒)。ONNX RuntimeとTensorRTは実験イメージに含まれていないため 未実測です(ダウンロードしていません)。 シリーズ「ビジョンモデルを読む3」第3回。精度削減については 第2回にあります。
目次
- 1. 直観:料理人と冷蔵庫
- 2. 実際にやってみる(ノートPCのCPU)
- 3. グラフコンパイラは何をするのか
- 4. ONNXとランタイム:モデルを持ち運ぶフォーマット
- 5. 自社での実測:ボールモデルへの torch.compile
- 6. 制限事項
- 7. 自社プロジェクトとのつながり
- 参考文献
1. 直観:料理人と冷蔵庫
GPUの計算装置は高速ですが、データはチップ外のメモリ(冷蔵庫)にあります。eager 実行は、「畳み込み → バッチ正規化 → ReLU → 加算」をそれぞれ別々に行います。畳み込みの結果を冷蔵庫に入れ、バッチ正規化が取り出して修正してまた入れ、ReLUがさらに取り出して入れます。大きな特徴マップであれば、計算よりも出し入れの時間の方が長くなります。 演算融合とは、「取り出したついでに4つすべてを行い、1度だけ書き戻す」ということです。
融合には2つの種類があります。
- 値を事前に畳み込む:推論時、バッチ正規化(BN)はチャネルごとに固定された乗算・加算であるため、直前の畳み込みの重みとバイアスにあらかじめ組み込むことができます。演算自体がなくなります。
- カーネルをまとめる:ReLU・加算・形式変換のように要素ごとに個別に行う演算を、前後のカーネルと一緒に1つのカーネルにまとめます。
BN 접기: y = γ·(W∗x + b − μ)/√(σ²+ε) + β
= (γ/√(σ²+ε))·W ∗ x + (γ·(b − μ)/√(σ²+ε) + β)
= W' ∗ x + b' ← 합성곱 하나로
2. 実際にやってみる(ノートPCのCPU)
code/03_fusion.py をノートPCのCPUスレッド1つで実行しました(7回の中央値)。
| 実験 | 結果 |
|---|---|
| 畳み込み(32→64、3×3、入力144×256)+BNの畳み込み | 最大差異4.77e-06(FP32の丸め誤差レベル)。時間5.80 → 5.80ms(差異なし) |
| 要素ごとのチェーン relu(x·1.5 + 0.5)·0.9、要素数1,677万個 | eager実行9.7ms → torch.compile 2.4ms(4.1倍)、結果の差異0 |
| 同じチェーンを手動でインプレース演算 | 7.4ms |
| 要素数102万個の加算を1,000回に分けて呼び出す | 0.13 → 0.88ms、呼び出しごとに約0.8µs |
- BNの畳み込みは数学的に同じ結果をもたらしましたが、このCPUでは時間が短縮されませんでした。大きな畳み込みの隣にあるBN1回の割合は小さいためです(最初の実行では2.6%減少し、2回目では0% — ノイズレベル)。
- 要素ごとのチェーンは異なります。eager 実行は大きな配列を4回読み書きしますが、コンパイルされたバージョンは1回読み込んで1回書き込みます。計算は同じでメモリの往復だけが減ったのに4倍高速化しました。このマシンでの最初のコンパイルは、キャッシュがない状態で5.4秒、キャッシュがある状態で0.6秒でした。
- 手動でインプレース演算を使用しても7.4msにとどまりました。新しい配列の割り当ては減らしたものの、依然として4回読み書きしています。 人がPythonレベルで行える最適化には限界があり、その下の1層(カーネル)を変更しなければ 大幅には短縮されません。
3. グラフコンパイラは何をするのか
TVM(Chenほか、2018) はモデルをグラフとして受け取り、演算を4つに分類します。要素ごとの処理(injective)、まとめ上げる処理(reduction)、出力に要素ごとの処理を付加できる複雑な処理(畳み込みのような complex-out-fusable)、そしてまとめられない処理です。ルールに従い、最初の3つを統合します。演算融合だけで1.2〜2倍高速化し(2015年デスクトップ級GPU、畳み込み+BN+ReLUなどの4つのタスク、3節・図4)、他のフレームワークと比較してエンドツーエンドで1.6〜3.8倍でした(6.1節)。
PyTorch 2(Anselほか、2024) の torch.compile は2つの部分から構成されます。
- TorchDynamo:Pythonバイトコードをインターセプトして演算グラフをキャプチャします。グラフ化できないPythonコードが現れた場合、その前後で分割し、残りはそのまま実行します。グラフキャプチャのオーバーヘッドは推論において約5%でした(表2、データセンターGPU)。
- TorchInductor:キャプチャしたグラフを融合し、GPU用のTritonカーネル(またはCPU用のC++)を生成します。
データセンターGPUにおいて、TIMM画像モデル62個の推論がeager実行よりもFP32で2.48倍、FP16で2.77倍高速になりました(表3)。融合とインライン化の両方を無効にすると、むしろ0.80倍となり、eager実行よりも遅くなりました(表4、HuggingFace FP16推論)。 速度の大部分は融合に由来することを意味します。冒頭の「推論2.27倍」がどの組み合わせの平均であるか本文で確認できなかったため、ここでは表3・4の値のみを使用します。
4. ONNXとランタイム:モデルを持ち運ぶフォーマット
torch.compile はPyTorch内部で動作します。サーバーにPythonとPyTorchを置きたくない場合や、ブラウザ・スマートフォンで実行する必要がある場合は、モデルを 標準グラフファイルとしてエクスポートする必要があります。
- ONNX:演算の種類とグラフを定義した公開フォーマットです。PyTorchからエクスポートすると、他のランタイムが読み込みます。
- ONNX Runtime:ONNXを実行するランタイムです。グラフの最適化を3つの段階で行います。基本(定数畳み込み、畳み込み+BNなどの重複排除)、拡張(GELUやアテンションなどの複雑な融合、CPU・CUDAなど)、配置(メモリレイアウトをCPUに合わせる)。デフォルト設定ではすべて有効になります(公式ドキュメント)。実際の計算は 実行プロバイダ(CPU、CUDA、TensorRT、Core ML、ブラウザのWebGPU・WASMなど)が実行します。同じファイルをさまざまなデバイスに持ち運べることこそがONNXの価値です(第7回)。
- TensorRT:NVIDIA GPU専用のオプティマイザです。レイヤーを融合し、カーネル候補を実際に実行して高速なものを選定し、FP16・INT8に低精度化します。結果はそのGPUに合わせた「エンジン」ファイルになるため、別のGPUに移行する場合は再作成する必要があります。かつてはキャリブレーションデータでINT8の範囲を決定する「キャリブレーター」(エントロピー、min-max)を使用していましたが、公式ドキュメントによるとTensorRT 10.1以降ではキャリブレーターはサポート終了となり、モデル内に量子化・逆量子化ノードを挿入する明示的な方式に変更されました。
5. 自社での実測:ボールモデルへの torch.compile
TT Scoreのボールモデル(BlurBall、3フレームウィンドウ、512×288、バッチ1)を実験用GPUで測定しました。
| 実行方式 | FP32 | FP16 | 初回呼び出し(コンパイル) |
|---|---|---|---|
| eager実行(現在の運用方式) | 15.09ms | 9.77ms | — |
| eager実行、cuDNN自動チューニング無効 | 15.08ms | 9.75ms | — |
| eager実行、channels_last(NHWC) | 18.62ms | 9.02ms | — |
| torch.compile デフォルト | 12.91ms(1.17倍) | 4.61ms(2.12倍) | 10〜11秒 |
| torch.compile reduce-overhead(CUDAグラフ) | 12.55ms | 4.34ms(2.25倍) | 7〜8秒 |
| torch.compile max-autotune(CUDAグラフなし) | 12.79ms | 4.48ms | 41〜46秒 |
(FP16はautocastを使用。倍率は同じ精度のeager実行と比較したもの。)
同じ実行をプロファイラーで確認したところ、eager実行のFP32の1回の実行で CUDAカーネル478個が呼び出され、GPUの稼働時間14.88msのうち畳み込みが8.09ms(72回)、バッチ正規化が3.27ms(71回、22%)でした。カーネル時間の合計がウォールクロック時間(15.09ms)とほぼ同じであるため、GPUは休止していませんでした。つまり、呼び出しコストではなく メモリを行き来する要素ごとの演算が問題であり、コンパイラがこれをまとめて排除したと解釈しています(コンパイルされたバージョンのカーネル数はカウントしていません)。
読み解き方のポイントです。
- FP16においてより大きな効果が得られました。 FP16のeager実行では畳み込みは高速化されたものの、BN・ReLU・形式変換がそのまま残り割合が大きくなっていましたが、コンパイラがまさにそれらを削除しました。第2回の精度削減と第3回の融合は掛け算の効果を生みます(15.09 → 4.34ms、3.5倍)。
- CUDAグラフ(reduce-overhead)によりさらにわずかに短縮されました。 1回のエクスポートにおけるカーネルをグラフとしてまとめて一度にGPUに渡すため、カーネルごとのPython・ドライバーの呼び出しコストがなくなります。このモデルではGPUがすでに高負荷であったため、利益は小さくなりました(4.61 → 4.34ms)。
- 時間をかけて選択させても速くなるとは限りませんでした。 max-autotuneは45秒を費やしたにもかかわらず、デフォルトよりも速くありませんでした。このGPUは行列積の自動チューニングを使用するには計算ユニットが少ないという警告も表示されました。
- メモリレイアウト(NHWC)はFP16でのみメリットがあり(9.77 → 9.02ms)、FP32ではデメリットとなりました(15.09 → 18.62ms)。cuDNNの自動チューニングによる違いはありませんでした。「この設定が速い」という一般論よりも、そのモデル・その精度で実際に測定する方が正確です。
6. 制限事項
- 初回呼び出しに時間がかかります。 デフォルトモードでも10秒かかります。リアルタイムで稼働するプロセスであれば起動時にウォームアップすることで隠蔽できますが、入力形状(バッチサイズ)が変更されると再コンパイルされる可能性があります。バッチ1・2・3すべてをウォームアップする必要があります(第6回)。コンパイルキャッシュをディスクに残して再起動時間を短縮することは検証していません(未実測)。
- 結果が同一であるかを個別に確認する必要があります。 融合は加算の順序を変更するため、末尾の桁が変わる可能性があります。ボールの位置が一致しているかどうかは今回測定していません(未実測)。第2回のFP16の比較と同様に、正解フレームで確認してから有効化する必要があります。
- ONNX Runtime・TensorRTは測定できませんでした。 実験画像には含まれておらず、今回の作業で新しいパッケージはダウンロードしていません。TensorRTのFP16・INT8エンジンがtorch.compileよりも速いかどうかについては、この記事では回答しません。
- 1つのGPU、1つのモデルの結果です。Transformerのように大きな行列積が中心のモデルでは、カーネル融合の恩恵が異なる場合があります。
7. 自プロジェクトへの適用
- 運用担当者の現在の実行環境はFP32です。第5節のとおりであれば、FP16+コンパイルによって1ウィンドウあたり15msから4〜5msになり、同じGPUでより多くのセッションを受け入れたり、他の作業と共有したりするための余裕が生まれます。有効化する前にすべきことは2つあります。録画された試合の正解フレームを使ってボールの位置や判定が一致しているかを確認すること(第2弾の方法)、そして本番運用GPUで再度測定を行うこと(GPUが変わればメリット・デメリットが入れ替わります — 第6弾)。
- 開発者は起動時にバッチサイズ1・2・3をウォームアップします。コンパイルを有効にすると、そのウォームアップがそのままコンパイル処理になるため、準備時間が10秒単位で増加します。準備完了の確認(
/healthz)はすでにウォームアップの後に開くため、ユーザーがその時間を目にすることはありません。 - スマホやブラウザで実行するモデルについては、まずONNX形式でエクスポートすることが先決です(第7弾)。
T.T.の一言。 同じ料理を冷蔵庫の開閉を減らして作ったら、2倍速くなりました。材料(重み)もレシピ(モデル)もそのままなのに!
参考文献
- Chen, T. et al. (2018). TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI 2018. https://arxiv.org/abs/1802.04799
- Ansel, J. et al. (2024). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. ASPLOS 2024. https://doi.org/10.1145/3620665.3640366
- ONNX. https://onnx.ai/ · ONNX Runtimeグラフ最適化. https://onnxruntime.ai/docs/performance/model-optimizations/graph-optimizations.html
- NVIDIA TensorRTドキュメント(How TensorRT Works, Explicit Quantization). https://docs.nvidia.com/deeplearning/tensorrt/latest/architecture/how-trt-works.html
- 独自の測定: 本シリーズの実測値(元データはリポジトリ外の研究記録)
- サンプルコードと出力: このフォルダの
code/03_fusion.py,code/03_fusion.out.txt