当前位置：首页 > news >正文

torch 中 model.eval() 和 model.train() 的作用

news 来源：原创 2025/6/7 4:55:12

在 PyTorch 中，model.train() 的作用是将模型切换到训练模式（training mode），主要影响模型中某些特定层（如 Dropout 和 BatchNorm）的行为，使其在训练时启用随机性和动态统计量。以下是详细说明：

核心功能

启用随机性层：：
- Dropout：在训练模式下，按设定的概率随机丢弃神经元（防止过拟合）。
- BatchNorm：使用当前 batch 的均值和方差进行归一化，并更新移动平均统计量（用于后续的评估模式）。
确保训练时的动态行为：训练模式下，模型的输出依赖于当前输入数据的随机性（如 Dropout）和动态统计量（如 BatchNorm），这对模型学习特征至关重要。

model.eval()的作用是将模型切换到评估模式（evaluation mode），主要影响模型中某些特定层（如Dropout和BatchNorm）的行为，使其在推理（测试）时表现一致且稳定。(实际上不使用dropout，model.eval()对这些不开dropout的大模型实际上没有影响。)：

核心功能

关闭随机性层：
- Dropout：在训练时随机丢弃神经元以防止过拟合，但在评估模式下会保留所有神经元。
- BatchNorm：在训练时使用当前batch的均值和方差进行归一化，并更新移动平均统计量；在评估模式下，则使用训练阶段累积的全局均值和方差，而非当前batch的数据。
确保输出稳定性：评估模式下，模型的输出仅依赖训练好的参数，避免因随机性（如Dropout）或统计量波动（如BatchNorm）导致测试结果不稳定。

`model.eval()`为什么需要配合`torch.no_grad()`？

model.eval()仅改变模型层的行为，而torch.no_grad()会禁用梯度计算，减少内存占用并加速推理。
通常在测试时同时使用两者：

推理阶段代码示例

model.eval()
with torch.no_grad():
    outputs = model(inputs)

Torch中模型训练评估模式演示代码

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# 定义模型（添加激活函数）
MyModel = nn.Sequential(
    nn.Linear(10, 20),
    nn.BatchNorm1d(20),# 训练时使用当前batch的统计量
    nn.ReLU(),
    nn.Dropout(0.5)     # 训练时随机丢弃50%的神经元
)

# 示例数据（假设是分类任务）
X_train = torch.randn(1000, 10)
y_train = torch.randint(0, 20, (1000,))
train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=64)

X_test = torch.randn(200, 10)
y_test = torch.randint(0, 20, (200,))
test_dataset = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=64)

# 定义损失函数和优化器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(MyModel.parameters(), lr=0.001)

model = MyModel

# 训练阶段
model.train()
for epoch in range(10):
    for data, targets in train_loader:
        optimizer.zero_grad()  # 重置梯度
        outputs = model(data)
        loss = loss_fn(outputs, targets)
        loss.backward()  #计算损失函数关于模型参数的梯度
        optimizer.step() #更新网络的权重和偏置等参数。


def calculate_accuracy(outputs, targets):
    preds = outputs.argmax(dim=1)
    correct = (preds == targets).sum().item()
    return correct / len(targets)

# 评估阶段
model.eval()
with torch.no_grad(): # 禁用梯度计算
    total_accuracy = 0
    for data, targets in test_loader:
        outputs = model(data)
        accuracy = calculate_accuracy(outputs, targets)
        total_accuracy += accuracy
    print(f"Test Accuracy: {total_accuracy / len(test_loader):.4f}")