PyTorch torch.no_grad 関数
torch.no_gradは、PyTorch で勾配計算を無効にするためのコンテキストマネージャーです。no_gradブロック内で作成されたテンソルは勾配を計算しないため、メモリ消費を大幅に削減し、推論速度を向上させることができます。
これはモデル推論(inference)および評価段階で不可欠であり、パフォーマンスを大幅に向上させ、メモリを節約できます。
関数定義
torch.no_grad()
パラメータ:
- パラメータはありません。これはコンテキストマネージャーです。
戻り値:
- コンテキストマネージャーを返します。そのコンテキスト内では勾配計算が無効になります。
使用例
例 1: 基本の使い方
例
import torch
# 勾配が必要なテンソルを作成
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# no_grad コンテキスト内で計算
with torch.no_grad():
y = x * 2
print("no_grad 内:", y.requires_grad)
# no_grad の外側
z = x * 2
print("no_grad 外:", z.requires_grad)
# 勾配が必要なテンソルを作成
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# no_grad コンテキスト内で計算
with torch.no_grad():
y = x * 2
print("no_grad 内:", y.requires_grad)
# no_grad の外側
z = x * 2
print("no_grad 外:", z.requires_grad)
出力結果は次のとおりです:
在 no_grad 中: False 在 no_grad 外: True
例 2: モデル推論
例
import torch
import torch.nn as nn
# 簡単なモデルを定義
model = nn.Linear(10, 2)
# 入力データを作成
x = torch.randn(1, 10)
# 推論時に no_grad を使用
with torch.no_grad():
output = model(x)
print("出力:", output)
# 勾配を計算する必要がない
# またはデコレータを使用
# @torch.no_grad()
# def predict(x):
# return model(x)
import torch.nn as nn
# 簡単なモデルを定義
model = nn.Linear(10, 2)
# 入力データを作成
x = torch.randn(1, 10)
# 推論時に no_grad を使用
with torch.no_grad():
output = model(x)
print("出力:", output)
# 勾配を計算する必要がない
# またはデコレータを使用
# @torch.no_grad()
# def predict(x):
# return model(x)
出力結果は次のとおりです:
输出: tensor([[0.0920, 0.3557]])
例 3: 評価モード
例
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 5),
nn.ReLU(),
nn.Linear(5, 2)
)
# 評価モードに切り替え
model.eval()
# テストデータを準備
test_input = torch.randn(5, 10)
# 評価時に勾配を無効化
with torch.no_grad():
predictions = model(test_input)
print("予測結果の形状:", predictions.shape)
import torch.nn as nn
model = nn.Sequential(
nn.Linear(10, 5),
nn.ReLU(),
nn.Linear(5, 2)
)
# 評価モードに切り替え
model.eval()
# テストデータを準備
test_input = torch.randn(5, 10)
# 評価時に勾配を無効化
with torch.no_grad():
predictions = model(test_input)
print("予測結果の形状:", predictions.shape)
出力結果は次のとおりです:
预测结果形状: torch.Size([5, 2])
例 4: メモリ使用量の比較
例
import torch
import torch.nn as nn
model = nn.Linear(1000, 1000)
# 大量の入力を作成
inputs = [torch.randn(100, 1000) for _ in range(100)]
# no_grad を使用しない(勾配履歴が記録される)
print("no_grad を使用しない場合:")
for inp in inputs[:5]:
_ = model(inp)
# no_grad を使用する(勾配履歴を記録しない)
print("no_grad を使用する場合:")
with torch.no_grad():
for inp in inputs[:5]:
_ = model(inp)
import torch.nn as nn
model = nn.Linear(1000, 1000)
# 大量の入力を作成
inputs = [torch.randn(100, 1000) for _ in range(100)]
# no_grad を使用しない(勾配履歴が記録される)
print("no_grad を使用しない場合:")
for inp in inputs[:5]:
_ = model(inp)
# no_grad を使用する(勾配履歴を記録しない)
print("no_grad を使用する場合:")
with torch.no_grad():
for inp in inputs[:5]:
_ = model(inp)
使用no_grad中間変数に勾配情報を保存する必要がないため、メモリ消費を大幅に削減できます。
関連関数
torch.enable_grad():勾配計算を有効にします(no_gradの逆)。torch.set_grad_enabled(grad):パラメータに応じて勾配計算を有効または無効にします。torch.inference_mode():より厳格な推論モードであり、勾配計算と autograd の両方を無効にします。
注意事項
- モデル推論と評価時は、常に使用してください。
torch.no_grad()メモリを節約し、速度を向上させるために。 - もし
no_gradブロック内で作成された変数をブロック外で使用する必要がある場合は、手動でコピーする必要があります。 - と
model.eval()組み合わせて使用すると効果的です。
その他の拡張
Pytorch torch リファレンスマニュアル