PyTorch torch.set_grad_enabled 関数

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

torch.set_grad_enabledこれは PyTorch で勾配計算を有効にするかどうかをグローバルに設定するための関数です。勾配計算を動的に有効または無効にでき、コンテキストマネージャーとは異なり、グローバルな状態を変更できる関数です。

これは、条件に応じて勾配計算を動的に制御する必要がある場合に非常に便利です。

関数定義

torch.set_grad_enabled(mode)

パラメータ:

mode(bool): もしTrueの場合は勾配計算を有効にし、もしFalseの場合は勾配計算を無効にします。

戻り値:

コンテキストマネージャーを返します。これはwith文で使用できます。

使用例

例1: 基本的な使い方

実例

import torch

# 勾配を必要とするテンソルを作成
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# 勾配計算を無効化
torch.set_grad_enabled(False)
y1 = x * 2
print("勾配を無効化した後:", y1.requires_grad)

# 勾配計算を有効化
torch.set_grad_enabled(True)
y2 = x * 2
print("勾配を有効化した後:", y2.requires_grad)

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

禁用梯度后: False
启用梯度后: True

例2: コンテキストマネージャーとして使用

実例

import torch

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

# コンテキストマネージャーを使用
with torch.set_grad_enabled(False):
y1 = x * 2
print("コンテキストマネージャー内:", y1.requires_grad)

y2 = x * 2
print("コンテキストマネージャー外:", y2.requires_grad)

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

在上下文管理器内: False
在上下文管理器外: True

例3: トレーニングと推論を動的に制御

実例

import torch
import torch.nn as nn

model = nn.Linear(10, 2)

def forward_pass(x, training=True):
"""training パラメータに応じて勾配を制御する"""
with torch.set_grad_enabled(training):
output = model(x)
print(f"training={training}, requires_grad={output.requires_grad}")
return output

# トレーニングモード
x = torch.randn(5, 10)
forward_pass(x, training=True)

# 推論モード
forward_pass(x, training=False)

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

training=True, requires_grad=True
training=False, requires_grad=False

例4: 勾配状態の保存と復元

実例

import torch

# 初期状態
print("初期状態:", torch.is_grad_enabled())

# 元の状態を返すコンテキストマネージャーを作成
old = torch.is_grad_enabled()

# 一時的に勾配を無効化
with torch.set_grad_enabled(False):
print("内部:", torch.is_grad_enabled())

# 自動的に復元される（ただし、ここでは手動で復元する必要がある）
torch.set_grad_enabled(old)
print("復元後:", torch.is_grad_enabled())

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

初始状态: True
在内部: False
恢复后: True

注意事項

set_grad_enabled関数として直接呼び出してグローバルな状態を変更することも、コンテキストマネージャーとして使用することもできます。
関数として呼び出す場合は、元の状態を手動で復元する必要があります。そうしないと、後続のコードに影響を及ぼします。
状態が正しく復元されるように、コンテキストマネージャー方式の使用を推奨します。
グローバルな状態への影響に注意し、複雑なコードでは慎重に使用してください。

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

その他の拡張

PyTorch torch.set_grad_enabled 関数

関数定義

使用例

例1: 基本的な使い方

実例

例2: コンテキストマネージャーとして使用

実例

例3: トレーニングと推論を動的に制御

実例

例4: 勾配状態の保存と復元

実例

関連関数

注意事項