AI LAB / KEYWORD DICTIONARY

基礎知識

torch.compile(PyTorchコンパイル機能)とは?

torch.compileは、PyTorchコードをJITコンパイルして最適化されたカーネルとして実行し、高速化を図るPyTorch 2.0以降の機能です。

仕組みと使いどころ

torch.compileは、Pythonコードを追跡してPyTorch演算のグラフを取得し、最適化されたカーネルへコンパイルします。Python関数には呼び出し形式またはデコレーター形式で適用でき、torch.nn.Moduleにも利用できます。既存コードへの変更を抑えながら、Pythonの処理負荷やGPUの読み書きを減らせる場合があります。最初の数回はコンパイル時間が加わるため、通常実行より遅くなることがあります。追跡しにくいコードではグラフブレークが発生し、未対応部分を通常実行してからコンパイルを再開するため、最適化の機会が失われます。高速化の程度はモデル構成やバッチサイズなどによって異なります。

torch.compileの活用例

たとえば、こんな場面で

複数のPyTorch演算で構成した関数やtorch.nn.Moduleにtorch.compileを適用し、同じ処理を繰り返す場面で実行時間を比較できます。初回のコンパイル時間と、その後の継続実行時間は分けて測定します。

使うときのポイント

PyTorch 2.0以降で利用し、初回実行のコンパイル時間を含めて評価します。期待した高速化が得られない場合は、グラフブレーク、モデル構成、バッチサイズを確認します。

参考資料

Introduction to torch.compile — PyTorch Tutorials 2.15.0+cu132 documentation (新しいタブで開きます)

用語の説明や活用例に関連する参考資料です。

READ & LEARN

用語を知ったら、活用を知る。

AI LABで、AIの動向や業務に取り入れるヒントを読む。

記事一覧を見る