PyTorch推論最適化ガイド

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

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

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向けに設計された低精度フォーマットである。FP32と同等の数値範囲を保ちつつ、演算を高速化する。詳細は「4. NVIDIA Tensor Coreを利用した高速行列演算」で説明する。

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

GPU世代とアーキテクチャ

NVIDIA GPUは世代ごとに異なる機能を持つ。最適化手法の選択には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形式で高速処理できるようになった

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()より高速である（計算グラフの追跡を無効化するため）
ビュー追跡とバージョンカウンタの更新を無効化する（テンソルのメモリ管理オーバーヘッドを削減する）
推論専用の最適化により、全体のオーバーヘッドを削減する

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

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

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

原理

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

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

メモリ使用量の削減: 同じメモリ容量でより大きなモデルやバッチサイズを扱える
メモリ帯域幅の軽減: データ転送量が半分になり、メモリボトルネックが緩和される
Tensor Core活用: 専用ハードウェアによる高速演算が可能になる

実装方法

# モデルを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)

BF16とFP16の選択:

BF16: Ampere世代以降で利用可能である。FP32と同じ数値範囲を持ち、数値安定性が高い（推奨）
FP16: Volta世代以降で利用可能である。メモリ効率は良いが、数値範囲が狭く、オーバーフロー/アンダーフローのリスクがある

注意事項:

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

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

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

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

原理

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ビットに削減したフォーマットである。Tensor Coreは低精度演算に特化した設計のため、仮数部を削減したTF32形式を使用することで高速演算が可能になる。これにより、FP32データ型の行列演算を高速に処理できる。

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

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

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

実装方法

import torch

# TF32を有効化する（Ampere世代以降で効果がある）
# 'highest': 完全なFP32精度を使用する（デフォルト）
# 'high': TF32またはBF16ベースのアルゴリズムを使用する（推奨）
# 'medium': BF16を使用する（高速だが精度が低い）
torch.set_float32_matmul_precision('high')

注意: torch.set_float32_matmul_precisionは行列乗算（matmul）に影響する設定である。cuDNNを使用する畳み込み演算には影響しない。畳み込み演算でのTF32を制御するには、torch.backends.cudnn.allow_tf32を使用する。

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

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

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

原理

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

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

実装方法

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系モデルの場合）
    """
    model = YourModel.from_pretrained(...)

    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)
            else:
                # Volta/Turing世代ではFP16を使用する
                model = model.half()
        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


@torch.inference_mode()
def inference(model, input_data):
    """
    推論を実行する。推論モード設定はモデルロード時ではなく、推論実行時に適用する。
    """
    if device.type == 'cuda':
        if model.dtype == torch.bfloat16:
            input_data = input_data.to(dtype=torch.bfloat16)
        elif model.dtype == torch.float16:
            input_data = input_data.half()

    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を有効化する

主要なポイント

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

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

よくある問題と解決方法

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

解決策:

一部の層のみFP32に戻す
Ampere世代以降のGPUを使用している場合は、BF16を試す（FP16使用時）
スケーリングを調整する

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

解決策:

PyTorchのバージョンを確認する（2.0以上が必要）
動的な処理を含むモデルでは使用を避ける
fullgraph=Falseオプションを試す

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

解決策:

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