with torch.no_grad 和 model.eval() 的区别详解

原创已于 2025-07-16 15:43:05 修改 · 728 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #算法 #经验分享 #笔记 #pytorch

于 2025-07-16 15:18:57 首次发布

AIBase 专栏收录该内容

9 篇文章

订阅专栏

在 PyTorch 开发过程中，with torch.no_grad() 和 model.eval() 是两个非常重要且常用的操作。很多初学者容易混淆它们的功能和使用场景。本文将详细介绍它们的作用、区别以及典型用法，帮助你更好地理解这两个关键操作。

1. 什么是 `torch.no_grad()`？

torch.no_grad() 是 PyTorch 提供的上下文管理器，用于关闭自动求导机制中的梯度计算。它能节省内存和计算资源，适合在推理（inference）阶段使用。

核心作用：

禁止 autograd 跟踪张量上的操作，不构建计算图。
结果张量的 requires_grad 自动变为 False。
使得计算更高效、节省显存。

示例代码：

import torch

x = torch.tensor([2.0], requires_grad=True)
with torch.no_grad():
    y = x * 3
print(y.requires_grad)  # False，梯度关闭

2. 什么是 `model.eval()`？

model.eval() 是 PyTorch 中 nn.Module 类的一个方法，用于将模型切换到评估模式（evaluation mode）。它主要影响模型中的某些特定层，如 Batch Normalization 和 Dropout，使它们在推理时表现出不同的行为。

核心作用：

关闭 Dropout 层，确保神经元不再随机丢弃。
Batch Normalization 层使用训练过程中计算得到的均值和方差，而不是当前批次的统计数据。
不影响梯度计算，仅改变层的行为。

示例代码：

model = MyModel()
model.train()  # 训练模式，默认状态
output_train = model(input_data)
model.eval()   # 评估模式
output_eval = model(input_data)

3. 两者的区别和联系

方面	`torch.no_grad()`	`model.eval()`
作用	禁止梯度计算，节省内存和计算资源	切换模型到评估模式，改变层行为
影响对象	所有在上下文中的张量操作	模型中 Dropout、BatchNorm 等层
是否关闭梯度计算	是	否
典型使用场景	推理阶段，或不需要梯度时	推理阶段，确保模型层表现一致

联系：
通常在推理阶段会同时使用两者，model.eval() 保证模型层行为正确，with torch.no_grad() 避免无用的梯度计算，从而提高效率。

4. 详细示例演示

下面给出一个完整的推理示例，演示如何结合 model.eval() 和 torch.no_grad() 使用：

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.dropout = nn.Dropout(p=0.5)
        self.bn = nn.BatchNorm1d(20)
        self.fc2 = nn.Linear(20, 2)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        print(f"是否记录梯度:{x.requires_grad}")
        x = self.dropout(x)
        x = self.bn(x)
        x = self.fc2(x)
        return x

# 初始化模型和数据
model = SimpleNet()
model.train()  # 训练模式
input_tensor = torch.randn(5, 10)
print("-"*50)
print("训练模式:")
output_train = model(input_tensor)
print("输出:", output_train)

print("-"*50)
print("评估模式:")
model.eval()
output_eval = model(input_tensor)
print(f"输出:{output_eval}")

print("-"*50)
print("关闭梯度的评估模式：")
with torch.no_grad():
    output_eval_no_grad = model(input_tensor)
print("输出:", output_eval_no_grad)

输出结果如下：
在这里插入图片描述
可以发现model.train()下和model.eval()下同样的输入得到不同的输出，这是因为Dropout和BN在不同模式下行为的不同。
关闭梯度计算并不影响评估模式输出的结果，只是起到关闭计算图提高推理速度和内存消耗的效果。

5. 实践建议总结

训练阶段
- 使用 model.train() 以启用 Dropout 和 BatchNorm 的训练行为。
- 不要使用 torch.no_grad()，否则会影响梯度计算。
推理阶段
- 使用 model.eval() 关闭 Dropout 并固定 BatchNorm 的统计数据。
- 使用 with torch.no_grad() 禁用梯度计算，提高效率，节省显存。
注意事项
- model.eval() 只影响模型的内部层行为，不影响是否计算梯度。
- torch.no_grad() 是全局禁用梯度计算的上下文，不依赖模型状态。