【Python】torch.no_grad()函数详解和示例
torch.no_grad() 是 PyTorch 中用于禁用梯度计算的上下文管理器。在使用此上下文管理器包裹的代码块内，不会对任何变量进行梯度的计算。这通常用于两种情况：
- 当我们只需要计算模型的前向传播，而不需要反向传播时（即，我们不需要计算参数的梯度）。
- 当我们想要暂时关闭梯度追踪，以防止内存消耗过大。
下面是一个使用 torch.no_grad() 的简单例子：
import torch
# 假设我们有一个模型和一些数据
model = torch.nn.Linear(10, 1)
inputs = torch.randn(1, 10)
# 使用 torch.no_grad() 来禁止梯度追踪
with torch.no_grad():
outputs = model(inputs)
# 在这个范围之外，模型参数仍然可以更新
# 例如，optimizer.step() 不会影响到带有 torch.no_grad(): 的计算在这个例子中，即使我们在 with 块内做了前向传播，也不会计算任何参数的梯度。这在我们只是想要进行前向传播而不进行反向传播时非常有用。
评论已关闭