🚀

torch.compileの効果検証

2025/01/14に公開

最近、仕事で使っている機械学習モデルもう少し早くできないかなーと思っていたところ「深層ニューラルネットワークの高速化 ML Systems」という書籍を発見しました。迷わず購入し早速読み始めてみると、序盤からさまざまな学びがあり大変勉強になっています。
この本の最初の方に、高速化のための各手法をさらっと解説するような章があります。その中でtorch.compileについて軽く触れられている箇所があり少し気になったので簡単な実験をしてみました。

検証内容

以下3つのモデルをHuggingFaceのmodelsから引っ張ってきてコンパイル前後の推論速度を比較しました。

  • T5(base)
  • Whisper(base)
  • DETR(resnet-50)

使用したスクリプトはこちらです。

動作環境

  • OS -> macOS Sequoia
  • CPU -> M2
  • python -> 3.12

検証結果

モデル パラメータサイズ コンパイルなし(秒) コンパイルあり(秒) 差分(秒) 改善率(%)
T5 223M 0.13438 0.13818 -0.00380 -2.8
Whisper 72.6M 0.88821 0.69877 0.18944 21.3
DETR 41.6M 0.70080 0.56652 0.13428 19.2

WhisperとDETRに関しては思っていたよりもしっかりと速度改善がされている印象で少し驚きました。
T5についてはわずかに推論速度が遅くなっていて、compileによる改善はあまり見込め無さそうです。

おまけ

DETRに関してはHuggingFaceのブログに比較検証の結果を見つけました。
T4を使って検証した際に26%速度改善できておりGPUの方が効果は大きそうです。

Discussion