PyTorch torch.inference_mode 関数

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

torch.inference_modeは PyTorch における推論モード用のコンテキストマネージャーです。これはtorch.no_gradより厳格で、勾配計算を無効化するだけでなく、autogradエンジンのすべての追跡機能も無効化します。

これはモデル推論時にはno_gradより効率的で、メモリ使用量をさらに削減し、推論速度を向上させることができます。

関数定義

torch.inference_mode(mode=True)

引数:

  • mode(bool, オプション): もしTrue（デフォルト値）の場合、推論モードを有効にします。もしFalseの場合、推論モードを終了します。デフォルトはTrue。

戻り値:

  • コンテキストマネージャーを返します。そのコンテキストでは勾配とautogradが無効になります。

使用例

例1: 基本的な使い方

例

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# inference_mode コンテキスト内
with torch.inference_mode():
    y = x * 2
    print("inference_mode 内:", y.requires_grad)

# no_grad コンテキスト内
with torch.no_grad():
    z = x * 2
    print("no_grad 内:", z.requires_grad)

出力結果は次のとおりです：

在 inference_mode 中: False
在 no_grad 中: False

例2: no_gradとinference_modeの比較

例

import torch

# テンソルを作成
x = torch.randn(100, 100)

# inference_mode 内
with torch.inference_mode():
    # 大量の計算を実行
    for _ in range(10):
        x = torch.mm(x, x)

    # 計算が完了しても、コンテキスト内のテンソルは逆伝播に使用できません
    result = x.sum()

    # 勾配を必要とするテンソルに変換できるか確認
    print("inference_mode 内:", result.is_leaf)

# no_grad 内で同じ計算を行う
x2 = torch.randn(100, 100)
with torch.no_grad():
    for _ in range(10):
        x2 = torch.mm(x2, x2)
    result2 = x2.sum()
    print("no_grad 内:", result2.is_leaf)

出力結果は次のとおりです：

在 inference_mode 中: False
在 no_grad 中: True

例3: モデル推論

例

import torch
import torch.nn as nn

# 簡単なモデルを定義
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

model.eval()

# 入力データを作成
x = torch.randn(100, 10)

# inference_mode を使用して推論を実行
with torch.inference_mode():
    output = model(x)
    print("出力形状:", output.shape)
    print("出力 requires_grad:", output.requires_grad)

# デコレータを使用することもできます
@torch.inference_mode()
def predict(x):
    return model(x)

result = predict(x)
print("デコレータ方式 - 出力形状:", result.shape)

出力結果は次のとおりです：

输出形状: torch.Size([100, 5])
输出 requires_grad: False
装饰器方式 - 输出形状: torch.Size([100, 5])

例4: メモリ最適化の比較

例

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(1000, 1000),
    nn.ReLU(),
    nn.Linear(1000, 1000),
    nn.ReLU(),
    nn.Linear(1000, 10)
)

# 異なるモードのメモリ使用量をテスト
x = torch.randn(50, 1000)

print("コンテキストマネージャーを使用しない場合:")
_ = model(x)

print("nno_grad を使用:")
with torch.no_grad():
    _ = model(x)

print("ninference_mode を使用:")
with torch.inference_mode():
    _ = model(x)

使用inference_modeにより、メモリをさらに最適化できます。autogradエンジンを完全に無効化するためです。

関連関数

  • torch.no_grad(): 勾配計算を無効化しますが、一部のautograd機能は保持します。
  • torch.enable_grad(): 勾配計算を有効にします。
  • torch.is_inference_mode_enabled(): 推論モードが有効かどうかを確認します。

注意事項

  • inference_mode 比 no_gradより厳格で、より多くの機能を無効化します。
  • にinference_mode内で作成されたテンソルは非リーフノードとしてマークされ、逆伝播に使用できません。
  • モデル推論と評価時には使用を推奨しますinference_mode最高のパフォーマンスを得るために。
  • inference_modeとはno_gradネストして使用できません。

Pytorch torch 参考手册Pytorch torch リファレンスマニュアル

その他の拡張