当构建神经网络使用到dropout层等时,网络的正向传播后反向传播神经元的系数会有所不同,因此需要用.eval()和.train()来指定模型方向。
net.train()
- 作用:将模型设置为训练模式。
- 影响:
- 启用 Dropout 层:Dropout 会随机丢弃一部分神经元(以一定概率),用于防止过拟合。
- 启用 Batch Normalization 的训练行为:BatchNorm 会使用当前批次的均值和方差来进行归一化,并更新其内部的运行均值和方差。
- 使用场景:
- 在训练模型时需要调用 net.train()(通常是默认状态)。
- 确保模型的所有层都处于训练模式。
net.eval()
- 作用:将模型设置为评估模式。
- 影响:
- 禁用 Dropout 层:Dropout 不会随机丢弃神经元,而是使用所有神经元的完整输出。
- 禁用 Batch Normalization 的训练行为:BatchNorm 会使用训练过程中记录的全局均值和方差,而不是当前批次的均值和方差。
- 使用场景:
- 在验证或测试模型时需要调用 net.eval()。
- 确保模型的所有层都处于评估模式,避免因 Dropout 或 BatchNorm 的行为导致结果不一致。
示例
model.train() # 切换到训练模式
for data, target in train_loader: optimizer.zero_grad() output = model(data) loss = loss_fn(output, target) loss.backward() optimizer.step()
model.eval() # 切换到评估模式
with torch.no_grad(): # 禁用梯度计算(加速推理并节省内存) for data, target in test_loader: output = model(data) # 计算验证或测试指标
torch.no_grad() 与 eval() 的区别
- torch.no_grad():禁用梯度计算,用于加速推理和节省内存,但不会改变模型的模式。
- eval():切换模型到评估模式,但不会禁用梯度计算。
- 通常在评估时,两者会结合使用
一个完整示例
import torch
import torch.nn as nn
import torch.optim as optim # 定义一个简单的模型
class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.fc = nn.Linear(10, 1) self.dropout = nn.Dropout(0.5) self.bn = nn.BatchNorm1d(10) def forward(self, x): x = self.bn(x) x = self.dropout(x) x = self.fc(x) return x # 初始化模型、损失函数和优化器
model = SimpleModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01) # 训练阶段
model.train() # 切换到训练模式
for epoch in range(5): for data, target in train_loader: # 假设 train_loader 已定义 optimizer.zero_grad() output = model(data) loss = loss_fn(output, target) loss.backward() optimizer.step() # 验证阶段
model.eval() # 切换到评估模式
with torch.no_grad(): # 禁用梯度计算 for data, target in val_loader: # 假设 val_loader 已定义 output = model(data) val_loss = loss_fn(output, target) print(f"Validation Loss: {val_loss.item()}")