import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 设备设置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

使用 torch.device 可使你的代码具有可移植性，并在不同硬件配置上高效运行。

2. 数据准备

我们将使用 torchvision 加载和预处理 FashionMNIST 数据集。torchvision.transforms 提供了方便的数据增强和归一化工具。

# 转换
transform = transforms.Compose([
    transforms.ToTensor(), # 将 PIL 图像或 NumPy ndarray 转换为张量，并缩放到 [0,1] 范围
    transforms.Normalize((0.5,), (0.5,)) # 使用均值 0.5 和标准差 0.5 进行归一化
])

# 加载 FashionMNIST 数据集
train_dataset = torchvision.datasets.FashionMNIST(
    root='./data', train=True, download=True, transform=transform
)
test_dataset = torchvision.datasets.FashionMNIST(
    root='./data', train=False, download=True, transform=transform
)

# 创建 DataLoaders
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

# 供参考，类别名称
classes = ('T恤/上衣', '裤子', '套头衫', '连衣裙', '外套',
           '凉鞋', '衬衫', '运动鞋', '包', '踝靴')

这里，transforms.ToTensor() 将图像转换为 PyTorch 张量，而 transforms.Normalize() 将像素值调整到特定范围，这通常有助于训练稳定性。DataLoader 处理批处理、混洗和并行数据加载。

3. 定义模型

让我们定义一个简单的 CNN。如果你已经学习了第2章，这个结构看起来会很熟悉。我们继承 nn.Module，并在 __init__ 中定义层，在 forward 方法中定义前向传播。

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1, stride=1), # 28x28x1 -> 28x28x16
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2), # 28x28x16 -> 14x14x16
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1, stride=1), # 14x14x16 -> 14x14x32
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)  # 14x14x32 -> 7x7x32
        )
        self.fc_layers = nn.Sequential(
            nn.Linear(32 * 7 * 7, 128), # 展平 7*7*32 特征图
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1) # 为全连接层展平输出
        x = self.fc_layers(x)
        return x

# 实例化模型并将其移动到设备
model = SimpleCNN(num_classes=len(classes)).to(device)
print(model)

该模型包含两个卷积层，接着是最大池化，然后是两个全连接层。这是一种标准但有效的图像分类模型结构，适用于 FashionMNIST 等任务。

4. 损失函数和优化器

接下来，我们选择一个损失函数和一个优化器。对于多类别分类，CrossEntropyLoss 是一个常见选择。对于优化器，Adam 是一种流行且有效的算法。

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

注意我们如何将 model.parameters() 传递给优化器。这会告诉优化器在训练期间需要更新哪些张量（我们模型的权重和偏置）。学习率 lr 是一个你可以调整的超参数。

5. 训练循环

这是 PyTorch 显式操作的优势所在。我们手动遍历 epoch 和批次。

def train_one_epoch(epoch_index, tb_writer=None):
    running_loss = 0.
    last_loss = 0.
    correct_predictions = 0
    total_samples = 0

    # 将模型设置为训练模式
    model.train(True) # or model.train()

    for i, data in enumerate(train_loader):
        # 每个数据实例都是一个输入 + 标签对
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device) # 将数据移动到配置的设备

        # 为每个批次清除梯度！
        optimizer.zero_grad()

        # 为当前批次进行预测
        outputs = model(inputs)

        # 计算损失及其梯度
        loss = criterion(outputs, labels)
        loss.backward() # 为每个 requires_grad=True 的参数 x 计算 dloss/dx

        # 调整学习权重
        optimizer.step() # 使用梯度 x.grad 更新 x 的值

        # 收集数据并报告
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total_samples += labels.size(0)
        correct_predictions += (predicted == labels).sum().item()

        if i % 100 == 99: # 每 100 个批次记录一次
            last_loss = running_loss / 100 # 每个批次的损失
            current_accuracy = 100 * correct_predictions / (batch_size * (i + 1)) # 近似准确率
            print(f'  Epoch {epoch_index + 1}, Batch {i + 1:5d} loss: {last_loss:.3f}, Accuracy: {current_accuracy:.2f}%')
            if tb_writer:
                tb_writer.add_scalar('Training Loss', last_loss, epoch_index * len(train_loader) + i)
                tb_writer.add_scalar('Training Accuracy', current_accuracy, epoch_index * len(train_loader) + i)
            running_loss = 0.

    epoch_accuracy = 100 * correct_predictions / total_samples
    return last_loss, epoch_accuracy # 返回最后一个批次的损失和当前 epoch 的准确率

让我们分解批次循环中的关键步骤：

inputs, labels = inputs.to(device), labels.to(device): 当前批次的数据被移动到 GPU（如果可用）。
optimizer.zero_grad(): 这非常重要。PyTorch 默认会累积梯度。如果你不在每个批次开始时清除它们，你将累积来自前一个批次的梯度，导致更新不正确。
outputs = model(inputs): 前向传播，输入数据流经网络以产生预测。
loss = criterion(outputs, labels): 损失函数将模型的预测 (outputs) 与真实标签 (labels) 进行比较。
loss.backward(): 这是 autograd 自动微分发生作用的地方。它计算损失对所有 requires_grad=True 的模型参数的梯度。
optimizer.step(): 优化器使用 backward() 调用中计算的梯度更新模型的参数。

以下图表说明了典型 PyTorch 训练循环中一个批次的操作流程：

训练循环中处理一个批次的操作流程。

6. 评估循环

在每个 epoch 之后，或在训练结束时，你都会想在单独的数据集（例如，验证集或测试集）上评估你的模型，以了解其泛化能力。评估循环与训练循环相似，但有主要区别：

def evaluate_model(loader):
    # 将模型设置为评估模式
    model.eval() # or model.train(False)

    running_vloss = 0.0
    correct_predictions = 0
    total_samples = 0

    # 禁用梯度计算以进行评估
    with torch.no_grad(): # 重要！
        for i, vdata in enumerate(loader):
            vinputs, vlabels = vdata
            vinputs, vlabels = vinputs.to(device), vlabels.to(device)

            voutputs = model(vinputs)
            vloss = criterion(voutputs, vlabels)
            running_vloss += vloss.item()

            _, predicted = torch.max(voutputs.data, 1)
            total_samples += vlabels.size(0)
            correct_predictions += (predicted == vlabels).sum().item()

    avg_vloss = running_vloss / (i + 1) # 所有验证批次的平均损失
    accuracy = 100 * correct_predictions / total_samples
    print(f'Validation Loss: {avg_vloss:.3f}, Validation Accuracy: {accuracy:.2f}%')
    return avg_vloss, accuracy

评估循环的主要区别：

model.eval(): 这会将模型设置为评估模式。这很重要，因为某些层（如 Dropout 和 BatchNorm）在训练和评估期间表现不同。model.eval() 确保它们使用评估时的行为。
with torch.no_grad():: 这个上下文管理器会禁用梯度计算。在评估期间，我们不需要计算梯度，这可以节省内存和计算。它确保评估过程的任何部分都不会意外修改模型的梯度。

7. 组合所有部分：主训练脚本

现在，让我们将这些函数组合成一个主脚本，用于训练模型几个 epoch 并进行评估。

num_epochs = 5 # 仅作演示；通常需要更多 epoch
train_losses = []
val_losses = []
train_accuracies = []
val_accuracies = []

print("Starting Training...")
for epoch in range(num_epochs):
    print(f'EPOCH {epoch + 1}:')

    # 训练
    model.train(True)
    avg_loss, train_acc = train_one_epoch(epoch) # 为简单起见，此处不使用 TensorBoard 写入器
    train_losses.append(avg_loss)
    train_accuracies.append(train_acc)

    # 验证
    model.eval()
    avg_vloss, val_acc = evaluate_model(test_loader)
    val_losses.append(avg_vloss)
    val_accuracies.append(val_acc)

    print(f'EPOCH {epoch + 1} Summary: Train Loss: {avg_loss:.3f}, Train Acc: {train_acc:.2f}%, Val Loss: {avg_vloss:.3f}, Val Acc: {val_acc:.2f}%')
    print("-" * 30)

print('Finished Training')

# 可选：保存模型
# torch.save(model.state_dict(), 'fashion_mnist_cnn.pth')
# print('模型已保存到 fashion_mnist_cnn.pth')

该脚本迭代 num_epochs 次。在每个 epoch 中，它会调用 train_one_epoch 在训练数据上训练模型，然后调用 evaluate_model 评估其在测试数据上的性能。我们还会收集损失和准确率值，这对于后续绘图很有用。

8. 训练进度可视化（可选）

一种常见做法是绘制训练和验证损失/准确率随 epoch 变化的图表，以监测过拟合情况并了解训练动态。

import matplotlib.pyplot as plt

epochs_range = range(1, num_epochs + 1)

plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_losses, label='Training Loss', color='#1c7ed6', marker='o')
plt.plot(epochs_range, val_losses, label='Validation Loss', color='#fd7e14', marker='x')
plt.title('Training and Validation Loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_accuracies, label='Training Accuracy', color='#1c7ed6', marker='o')
plt.plot(epochs_range, val_accuracies, label='Validation Accuracy', color='#fd7e14', marker='x')
plt.title('Training and Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy (%)')
plt.legend()

plt.tight_layout()
plt.show()

如果你运行完整脚本，应该会看到类似的图表（实际值会因运行而异）：

显示训练和验证损失随 epoch 减少的典型趋势的示例图。

给 TensorFlow 开发者的建议

对于 TensorFlow Keras 的开发者而言，model.compile() 和 model.fit() 抽象了许多细节，PyTorch 的做法可能最初看起来更冗长。然而，这种显式操作是 PyTorch 的优势之一。

完全控制: 你可以直接控制训练过程的每一步。这使得实现自定义训练逻辑、非标准反向传播方案或复杂指标计算变得更容易。
透明性: 你可以清楚地看到发生了什么：梯度何时清零，损失如何计算，权重何时更新。这对于调试和理解模型行为非常有帮助。
灵活性: 需要在优化器步骤之前修改梯度吗？想对模型的不同部分应用不同的优化器吗？PyTorch 的手动循环使这些任务变得简单。

虽然 Keras 为标准工作流程提供了便利，但掌握 PyTorch 训练循环有助于你处理更广泛的研究问题并实现更复杂的训练程序。这个动手练习已经打下了基础。随着你的进步，你会发现自己会为这个基本循环添加更多功能，例如学习率调度器、早期停止和自定义日志记录，所有这些都自然地融入到这种显式结构中。

这部分内容有帮助吗？