如何在 PyTorch 中训练模型

在PyTorch中训练模型，需定义模型、损失函数和优化器，加载并预处理数据，通过多个训练轮次（epoch）进行前向传播、计算损失、反向传播和参数更新。此外，还需评估模型性能，并可视化训练过程。PyTorch的灵活性和强大功能使其成为训练深度学习模型的优选工具。

在 PyTorch 中训练模型的典型流程可分为以下步骤，附上代码示例和说明：

1. 准备数据

使用 Dataset 和 DataLoader 加载数据：

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

# 示例：创建随机数据（假设输入维度为 32，输出为 10 类）
X = torch.randn(1000, 32)  # 1000 个样本
y = torch.randint(0, 10, (1000,))

# 封装为 Dataset 和 DataLoader
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

2. 定义模型

继承 nn.Module 定义网络结构：

import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(32, 64)    # 输入层到隐藏层
        self.fc2 = nn.Linear(64, 10)    # 隐藏层到输出层

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleModel()

3. 定义损失函数和优化器

criterion = nn.CrossEntropyLoss()        # 分类任务常用交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # 优化器选择 Adam

4. 训练循环

核心步骤：前向传播、计算损失、反向传播、参数更新

num_epochs = 10

for epoch in range(num_epochs):
    model.train()  # 确保模型处于训练模式（影响 Dropout/BatchNorm 等层）

    for batch_X, batch_y in dataloader:
        # 前向传播
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)

        # 反向传播与优化
        optimizer.zero_grad()  # 清空梯度（必须的步骤！）
        loss.backward()        # 计算梯度
        optimizer.step()       # 更新参数

    # 可选：打印每个 epoch 的损失
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

5. 可选：保存模型

torch.save(model.state_dict(), "model_weights.pth")

关键细节说明

设备选择：用 model.to(device) 和 batch_X.to(device) 将数据/模型移到 GPU（若有）：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleModel().to(device)

梯度清零：每次反向传播前需调用 optimizer.zero_grad()，避免梯度累积。
数据预处理：真实场景中需对数据做归一化（如 transforms.Normalize）或使用预定义数据集（如 torchvision.datasets.MNIST）。
验证步骤：通常在每个 epoch 后评估验证集性能，需调用 model.eval() 并禁用梯度计算：
```
with torch.no_grad():
   # 计算验证集精度/损失
```