PyTorch推論最適化ガイド

【概要】本文書は、PyTorchを使用した深層学習モデルの推論速度向上を目的とした最適化手法について解説する。効果はモデル構造、GPU世代、入力サイズによって異なるため、実際の環境での検証が必要である。

【本文書の前提と読み方】

スコープ: 本文書は推論時の最適化のみを対象とする。
「混合精度」の用語整理: 本文書では、モデル全体を低精度（FP16/BF16）に変換する手法を「混合精度推論」として扱う。一方、PyTorchが標準で提供する狭義の混合精度API（torch.autocast）は、演算ごとに精度を切り替え、数値安定性が必要な演算のみFP32を維持する仕組みである。両者の差異は「2. 混合精度推論」で補足する。
各最適化の併用順序: 複数の手法を併用する際は、「精度変換（FP16/BF16/TF32）→ torch.compile」の順で適用する。コンパイル後に精度を変更すると、コンパイル結果が無効化される場合があるためである。
効果量の検証とウォームアップ: torch.compileの初回コンパイルやcudnn.benchmarkのアルゴリズム探索には時間を要する。これらの初回コストを除外して効果を測るため、ベンチマーク計測前に数回のウォームアップ実行を行うこと。
動作確認バージョン: 本文書はPyTorch 2.x系を前提とする。特にTF32制御のAPIはPyTorch 2.9で刷新されたため、使用中のバージョンの公式ドキュメントを併せて確認することを推奨する。

関連する外部サイトを以下に示す。

TF32に関する詳細: PyTorch CUDA Semantics
torch.compileに関する詳細: torch.compile Tutorial
混合精度に関する詳細: Automatic Mixed Precision
torch.set_float32_matmul_precisionに関する詳細: 公式ドキュメント
cuDNNに関する詳細: NVIDIA cuDNN Documentation

用語リスト

本章では、最適化手法を理解するために必要な用語を解説する。浮動小数点数の表現とGPUアーキテクチャの知識は、混合精度推論やTF32の理解に必要となる。

浮動小数点数の表現

浮動小数点数は「符号部」「指数部」「仮数部」の3つの要素で構成される。

仮数部: 数値の精度を決定する部分である。ビット数が多いほど高精度となる。
指数部: 数値の範囲（表現可能な最大値と最小値の幅）を決定する部分である。

各フォーマットの仮数部ビット数は、FP32が23ビット、FP16が10ビット、BF16が7ビット、TF32が10ビットである。

※ IEEE 754規格（浮動小数点数の国際標準規格）では仮数部に暗黙の先頭1ビットがあるため、実質的な精度はフィールドより1ビット多い（例: FP32は24ビット精度）。本文書では便宜上、フィールドのビット数を記載している。

TF32（TensorFloat-32）: Ampere世代以降のTensor Core向けに設計された演算モードである。入力を指数部8ビット・仮数部10ビットに丸めて計算し、結果はFP32精度で累算する。これにより、FP32と同等の数値範囲を保ちつつ演算を高速化する。詳細は「4. NVIDIA Tensor Coreを利用した高速行列演算」で説明する。

BF16（Brain Float 16）: Googleが開発したフォーマットである。FP32と同じ指数部（8ビット）を持ち、仮数部を7ビットに削減している。FP16と比較して数値範囲が広く、深層学習での数値安定性が高い。

GPU世代とアーキテクチャ

NVIDIA GPUは世代ごとに異なる機能を持つ。最適化手法の選択にはGPU世代の把握が必要である。

Compute Capability: NVIDIAがGPUの機能レベルを示すために定めた番号である。数値が大きいほど新しい機能をサポートする。

Volta世代（2017年〜）: 代表的なGPUはV100である。Compute Capability 7.0。この世代でTensor Coreが初めて搭載され、FP16による高速行列演算が可能になった。
Turing世代（2018年〜）: 代表的なGPUはT4、RTX 20xxシリーズである。Compute Capability 7.5。FP16 Tensor Coreをサポートする。
Ampere世代（2020年〜）: 代表的なGPUはA100、RTX 30xxシリーズである。Compute Capability 8.0以上。TF32とBF16に対応したTensor Coreを搭載し、FP32データ型を使用した演算をTF32形式で高速処理できる。
Ada Lovelace / Hopper世代以降（2022年〜）: 代表的なGPUはRTX 40xxシリーズ、H100である。Compute Capability 8.9以上。Ampere世代の機能に加え、FP8など、より低精度のフォーマットによる高速演算に対応する。後継のBlackwell世代（2024年〜、B100/B200など）も同様の方向で機能を拡張している。本文書の手法はこれらの世代でもAmpere世代と同様に適用できる。

Tensor Coreとは

Tensor Coreは、NVIDIA GPUに搭載された行列演算専用ハードウェアである。従来の演算ユニット（CUDA Core）はFP32の精度で汎用的な演算を行うが、Tensor Coreは低精度フォーマット（TF32、FP16、BF16など）の行列演算に特化している。精度を制限することで同じシリコン面積に多くの演算器を搭載でき、行列演算のスループット（単位時間あたりの処理量）が向上する。

PyTorchにおけるTensor

PyTorchの「Tensor」は多次元配列のデータ構造を指し、数学的なテンソル（多次元配列）をプログラムで扱えるようにしたものである。GPU上で効率的に計算できる。

勾配計算

勾配計算とは、ニューラルネットワークの学習において、損失関数の各パラメータに対する微分（勾配）を計算する処理である。推論時には不要だが、学習時には必須となる。

その他の用語

オーバーヘッド: 本来の処理以外に発生する付加的な処理時間やメモリ消費のことである。
GPUカーネル: GPU上で実行される関数の単位である。カーネルの起動には一定のオーバーヘッドが伴う。
eager mode: PyTorchの標準的な実行モードである。演算を定義と同時に実行する方式を指す。

最適化手法

本章では5つの最適化手法を紹介する。各手法は独立して適用可能だが、組み合わせることで効果が増加する場合がある。以下では、適用範囲が広く基本的な手法から順に説明する（順序は厳密な適用範囲の広狭ではなく、理解しやすさを優先している）。

各手法の関係を以下に示す。

推論モード設定: すべての推論処理で使用する基本設定である。
混合精度推論とTF32: いずれもTensor Coreを活用する手法である。混合精度推論を使用する場合、TF32の効果は限定的となる。
torch.compile: 他の最適化と併用可能である。精度変換後に適用する。
cuDNN最適化: CNN系モデル向けの最適化である。他の手法と併用可能である。

1. 推論モード設定

対象: 推論時（勾配計算が不要な場合）

原理

学習時、PyTorchは自動微分のために計算グラフ（演算履歴）を保持する。torch.inference_mode()は、推論時に不要な計算グラフの構築を無効化し、メモリとCPUオーバーヘッドを削減する。

実装方法

# デコレータとして使用する
@torch.inference_mode()
def predict(model, data):
    return model(data)

# またはコンテキストマネージャーとして使用する
with torch.inference_mode():
    output = model(input_data)

効果:

torch.inference_mode()はtorch.no_grad()より高速である。計算グラフの追跡を無効化するためである。
ビュー追跡とバージョンカウンタの更新を無効化し、テンソルのメモリ管理オーバーヘッドを削減する。

注意: torch.inference_mode()下で生成されたテンソルは、後からautograd計算に再投入できない制約がある。推論結果を学習ループに戻すような用途ではtorch.no_grad()を使用する。

適用シーン: すべての推論処理で使用する。学習時に使用すると勾配計算ができなくなるため、学習と推論を同一コード内で行う場合は適用範囲に注意が必要である。

2. 混合精度推論（FP16/BF16）

対象: Tensor Coreを持つGPU（Volta世代以降）

原理

FP16（16ビット浮動小数点）やBF16（Brain Float 16）は、FP32より少ないビット数で数値を表現する。メモリ使用量が半分になり、メモリ帯域幅の制約が軽減される。さらにTensor Coreによる高速演算が可能になる。

「混合精度」と呼ばれる理由は、本来、モデル全体を低精度に変換しつつ、数値安定性が求められる一部の演算でFP32を維持する手法を指すためである。

混合精度推論を使用すると、以下の3つの効果が得られる。

メモリ使用量の削減: 同じメモリ容量でより大きなモデルやバッチサイズを扱える（推論時に固定バッチで運用する場合でも、常駐メモリの削減により他プロセスとの共存が容易になる）。
メモリ帯域幅の軽減: データ転送量が半分になり、メモリボトルネックが緩和される。
Tensor Core活用: 専用ハードウェアによる高速演算が可能になる。

広義の混合精度と狭義の混合精度: PyTorchで「混合精度」と言う場合、以下の2通りの実現手段がある。

広義（モデル全体を低精度に変換）: model.half()やmodel.to(dtype=torch.bfloat16)で全パラメータと演算を一律に低精度化する。実装が簡潔だが、数値安定性に問題が出る場合は手動で一部の層をFP32に戻す必要がある。本文書のコード例ではこちらを採用する。
狭義（torch.autocastによる自動切り替え）: 演算ごとに精度を自動で切り替え、数値安定性が必要な演算（softmax、損失計算等）はFP32を維持する。PyTorch標準APIであり、より安全である。詳細は公式ドキュメントを参照のこと。

実装方法

# --- 広義の混合精度: モデル全体を一律に低精度化する ---
# モデルをFP16に変換する
model = model.half()
input_data = input_data.half()

# またはBF16に変換する（Ampere世代以降）
model = model.to(dtype=torch.bfloat16)
input_data = input_data.to(dtype=torch.bfloat16)

# 推論を実行する
with torch.inference_mode():
    output = model(input_data)

# --- 狭義の混合精度: torch.autocastを使用する場合の例 ---
# モデルはFP32のまま保持し、autocast内で自動的に精度が切り替わる
# with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.bfloat16):
#     output = model(input_data)

BF16とFP16の選択:

BF16: Ampere世代以降で利用可能である。FP32と同じ数値範囲を持ち、数値安定性が高いため推奨される。
FP16: Volta世代以降で利用可能である。メモリ効率は良いが、数値範囲が狭く、オーバーフローやアンダーフロー（数値が表現可能な範囲を超えて、無限大や0として扱われる現象）のリスクがある。

注意事項:

一部のモデルや層はFP32精度が必要な場合がある。
Turing世代（T4など）ではBF16はサポートされない。FP16のみ利用可能である。

適用シーン: 大規模モデル（GPT、BERT、Stable Diffusionなど）で効果的である。小規模モデルでは効果が限定的な場合がある。

3. torch.compileによるモデル全体のコンパイル最適化

対象: PyTorch 2.0以上。効果はモデルの構造に依存し、動的処理が多いモデルでは効果が薄い。

適用順序の注意: 精度変換（FP16/BF16/TF32等）はtorch.compile適用前に行うこと。コンパイル後に精度を変更すると、コンパイル結果が無効化され再コンパイルが発生する場合がある。

原理

PyTorchは通常、各演算を逐次的に実行する（eager mode）。torch.compileは実行前にモデル全体を解析し、以下の最適化を行う。

演算の融合: 複数の小さな演算（例: 行列演算、活性化関数、正規化）を1つのGPUカーネルにまとめ、中間結果のメモリ書き込みと読み込みを削減する。
メモリアクセスの最適化: データの配置を工夫し、GPUメモリへのアクセス回数や帯域幅の使用を削減する。
不要な演算の削除: 結果に影響しない演算を事前に除外する。

これにより、GPUの実行効率が向上する。

実装方法

# PyTorch 2.0以上で利用可能である
if hasattr(torch, 'compile'):
    model = torch.compile(model)

注意事項:

初回実行時にコンパイル時間が発生する。所要時間はモデルサイズに依存する。ベンチマーク計測時は数回のウォームアップ実行を行うこと。
mode='max-autotune'を指定すると、さらに高速になる場合があるが、初回コンパイル時間が大幅に増加する。多くのケースでは既定モードとの差が小さいため、まず既定モードで効果を確認してから検討するとよい。

適用シーン: 静的な構造を持つモデル（ResNet、EfficientNetなど）で効果的である。動的な分岐を持つモデル（RNNの可変長処理など）では効果が限定的である。

4. NVIDIA Tensor Coreを利用した高速行列演算

対象: PyTorch 1.12以上、Ampere世代以降のCUDA GPU（RTX 30xxシリーズ以降など）

原理

TF32（TensorFloat-32）は、FP32の指数部（8ビット）を保持しつつ、入力を仮数部10ビットに丸めて計算し、結果はFP32精度で累算する演算モードである。Tensor Coreは低精度演算に特化した設計のため、仮数部を削減したTF32形式を使用することで高速演算が可能になる。これにより、FP32データ型の行列演算を高速に処理できる。

仮数部の削減により多少の精度低下がある。例えば、FP32では1.234567890が表現できるが、TF32では1.2345程度の精度となる。深層学習のような統計的処理では、この精度で十分な場合が多い。

トレードオフ: 計算速度の向上と仮数部精度の低下。

注意（デフォルト挙動）: PyTorch 1.12以降、行列乗算（matmul）におけるTF32はデフォルトで無効である。一方、cuDNNを使用する畳み込み演算では、TF32はデフォルトで有効である。これは、深層学習以外で数値計算の厳密性を要求するアプリケーションに配慮したためである。

注意（API刷新）: PyTorch 2.9以降、TF32をより細かく制御する新しいAPI（torch.backends.fp32_precision、torch.backends.cuda.matmul.fp32_precision、torch.backends.cudnn.conv.fp32_precisionなど）が導入された。これらは"ieee"（FP32精度）または"tf32"（TF32を許可）を指定する形式で、バックエンドや演算ごとの制御が可能である。従来のallow_tf32フラグは将来的に非推奨（deprecated）となる予定であり、新旧の設定の混在はサポートされない。本文書では広く普及しているtorch.set_float32_matmul_precisionを中心に説明するが、PyTorch 2.9以降を使用する場合は新APIの利用を検討すること。

実装方法

import torch

# TF32を有効化する（Ampere世代以降で効果がある）
# 'highest': 完全なFP32精度を使用する（デフォルト）
# 'high'   : 行列乗算でAmpere世代以降のGPUを使用している場合はTF32が選択される。
#            一部の演算ではBF16ベースの近似アルゴリズムが選ばれることもある（推奨）
# 'medium' : より積極的にBF16ベースの近似アルゴリズムを使用する（高速だが精度が低い）
torch.set_float32_matmul_precision('high')

注意（適用範囲）: torch.set_float32_matmul_precisionは主に行列乗算（matmul）の内部精度に影響する設定である。cuDNNを使用する畳み込み演算でのTF32を制御するには、torch.backends.cudnn.allow_tf32（またはPyTorch 2.9以降のtorch.backends.cudnn.conv.fp32_precision）を使用する。

適用シーン: 大規模な行列演算が多いモデル（Transformer、大規模な全結合層など）で効果的である。混合精度推論を使用する場合は、主要な演算が低精度で行われるため、TF32の効果は限定的となる。

5. cuDNN最適化による畳み込み演算高速化

対象: GPU使用時、cuDNNインストール済み環境、入力サイズが固定のCNN系モデル

原理

cuDNN（CUDA Deep Neural Network library）は、NVIDIAが提供する深層学習演算ライブラリである。cudnn.benchmarkは、最初の実行時に複数の畳み込みアルゴリズムを試行し、最速のものを自動選択する。

トレードオフ: 初回実行の探索時間と以降の実行速度向上。

実装方法

if device.type == 'cuda':
    torch.backends.cudnn.enabled = True      # cuDNNを有効化する（デフォルトで有効）
    torch.backends.cudnn.benchmark = True    # 最適なアルゴリズムを自動選択する

注意: 入力サイズが変動する場合は逆効果となる。毎回アルゴリズム探索が実行され、処理が遅くなる。入力サイズが固定であることを確認してから有効化する。ベンチマーク計測時は数回のウォームアップ実行を行ってから計測すること。

適用シーン: CNN系のモデル（ResNet、EfficientNet、YOLOなど）で、入力サイズが固定の場合に効果的である。

適用タイミング

モデルロード後、推論開始前に最適化を適用する。以下は、これらの最適化手法を統合した実装例である。本文で説明した順序とは異なり、実装上の依存関係に基づいた順序で記述している。具体的には、精度変換をtorch.compileより前に行う必要がある。

def load_model(use_mixed_precision=True, fixed_input_size=True):
    """
    モデルをロードし、最適化を適用する。

    Args:
        use_mixed_precision: 混合精度を使用するかどうか
        fixed_input_size: 入力サイズが固定かどうか（CNN系モデルの場合）

    Returns:
        (model, model_dtype) のタプル。model_dtypeは推論時の入力変換に使用する。
    """
    model = YourModel.from_pretrained(...)
    model_dtype = torch.float32  # デフォルト

    if device.type == 'cuda':
        # cuDNN最適化（CNN系モデルで入力サイズ固定の場合のみ有効化する）
        torch.backends.cudnn.enabled = True
        if fixed_input_size:
            torch.backends.cudnn.benchmark = True

        # 混合精度またはTF32を選択する
        if use_mixed_precision:
            if torch.cuda.get_device_capability()[0] >= 8:
                # Ampere世代以降ではBF16を使用する
                model = model.to(dtype=torch.bfloat16)
                model_dtype = torch.bfloat16
            else:
                # Volta/Turing世代ではFP16を使用する
                model = model.half()
                model_dtype = torch.float16
        else:
            # 混合精度を使用しない場合はTF32を有効化する
            # 混合精度使用時も一部のFP32演算にTF32が適用される可能性があるが、
            # 主要な演算が低精度で行われるため効果は限定的である
            torch.set_float32_matmul_precision('high')

        # torch.compile最適化（PyTorch 2.0以上で利用可能）
        # 精度変換後にcompileすることで、変換後の精度に基づいた最適化が行われる
        if hasattr(torch, 'compile'):
            model = torch.compile(model)

    return model, model_dtype


@torch.inference_mode()
def inference(model, input_data, model_dtype=torch.float32):
    """
    推論を実行する。推論モード設定はモデルロード時ではなく、推論実行時に適用する。

    Args:
        model: load_model()でロード・最適化済みのモデル
        input_data: 入力テンソル
        model_dtype: load_model()が返したdtype。入力テンソルの精度合わせに使用する
                     （torch.compileでラップ後のモデルにはdtype属性が無いため、
                       外部から明示的に渡す設計としている）。
    """
    if device.type == 'cuda' and model_dtype != torch.float32:
        input_data = input_data.to(dtype=model_dtype)

    return model(input_data)

最適化手法の選択ガイド

実際の適用では、以下の判断基準に従って最適化手法を選択する。

GPU世代別の推奨設定

GPU世代	推奨される最適化	備考
Volta (V100)	inference_mode + FP16混合精度	TF32非対応、BF16非対応
Turing (T4, RTX 20xx)	inference_mode + FP16混合精度	TF32非対応、BF16非対応
Ampere (A100, RTX 30xx)	inference_mode + BF16混合精度 + torch.compile	混合精度を使用しない場合はTF32を有効化する
Ada / Hopper以降 (RTX 40xx, H100など)	inference_mode + BF16混合精度 + torch.compile	Ampereと同様。加えてFP8など、より低精度の手法も利用可能（本文書では扱わない）

※ torch.compileはPyTorch 2.0以上であればVolta/Turing世代でも利用可能である。動的構造を持たない静的なモデルでは、これらの世代でも効果が期待できる。表ではAmpere世代以降での代表的な組み合わせを示すにとどめている。

主要なポイント

GPU世代に応じた最適化手法の選択: Ampere世代以降ではBF16を、Volta/Turing世代ではFP16を使用する。TF32はAmpere世代以降で混合精度を使用しない場合に有効である。
混合精度（BF16/FP16）の効果: メモリ、帯域幅、演算速度の3つの面で効果がある。
複数の最適化の組み合わせ: 組み合わせることで効果が増加する場合がある。組み合わせ方には注意が必要である。例えば、混合精度使用時はTF32の効果が限定的となる。
環境での検証: 効果はモデルとハードウェアに依存するため、実際の環境での検証が必要である。本文書のガイドラインは一般的な傾向であり、個別のケースでは異なる結果となる場合がある。

付録: トラブルシューティング

よくある問題と解決方法

問題1: 混合精度で精度が低下する

解決策:

FP16使用時にオーバーフローやアンダーフローが疑われる場合は、Ampere世代以降のGPUでBF16へ切り替える（BF16はFP32と同じ数値範囲を持つため数値安定性が高い）。
数値安定性が必要な層を手動でFP32に戻す（例: layer.to(torch.float32)）。
モデル全体を低精度化する代わりに、torch.autocastによる狭義の混合精度に切り替え、数値安定性が必要な演算をPyTorchに自動でFP32維持させる。

問題2: torch.compileでエラーが発生する

解決策:

PyTorchのバージョンを確認する。2.0以上が必要である。
動的な処理を含むモデルでは使用を避ける。
fullgraph=Trueを指定している場合は、これを外す（デフォルトはfullgraph=False）。fullgraph=Trueはモデル全体が単一のグラフとしてコンパイル可能でなければエラーを発生させる設定である。fullgraph=Falseでは、コンパイルできない部分があってもエラーにせず、その部分はeager modeで実行される。

問題3: cuDNN benchmarkを有効にしたら遅くなった

解決策:

入力サイズが変動していないか確認する。
入力サイズが変動する場合はtorch.backends.cudnn.benchmark = Falseに設定する。

問題4: TF32の設定が効かない、または非推奨の警告が出る

解決策:

使用GPUがAmpere世代以降（Compute Capability 8.0以上）であることを確認する。Volta/Turing世代ではTF32は利用できない。
PyTorch 2.9以降では従来のallow_tf32フラグが将来的に非推奨となる。警告が出る場合は新API（torch.backends.cuda.matmul.fp32_precision = "tf32"、torch.backends.cudnn.conv.fp32_precision = "tf32"など）への移行を検討する。新旧の設定を混在させないこと。