pytorch训练流程 | 简单示例

MNIST 手写数字识别模型的训练过程。

在训练循环中，使用 交叉熵损失函数 和 随机梯度下降（SGD）优化器 进行模型训练。

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms

"""
	# 定义模型
	在 SimpleModel 类中定义了一个简单的神经网络模型，该模型包括一个全连接层，用于将输入的图像数据转换为数字标签的预测结果。
"""
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

"""
	# 准备数据
	准备数据：使用 torchvision 库加载 MNIST 数据集，并对数据进行预处理，包括转换为 Tensor 格式和归一化处理。
	然后通过 DataLoader 创建训练数据加载器，以便在训练过程中批量加载数据。
"""
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)


"""
	# 初始化模型和优化器
	创建模型实例 SimpleModel()，并定义优化器（这里使用随机梯度下降优化器）和损失函数（交叉熵损失）。
"""
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()


"""
	# 训练循环
	在每个 epoch 中，模型被设置为训练模式，然后遍历训练数据加载器中的每个批次。
	
	对于每个批次，首先将优化器的梯度清零，然后通过前向传播计算模型的预测输出，
	接着计算预测结果与真实标签之间的损失，通过反向传播更新模型参数。
	
	最后，打印每个 epoch 中损失的平均值，以及每个批次中损失的变化情况。
"""
epochs = 5
for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        # 优化器的梯度清零
        optimizer.zero_grad()
        # 前向传播计算模型的预测输出
        outputs = model(inputs)
        # 计算预测结果与真实标签之间的损失
        loss = criterion(outputs, labels)
        # 反向传播更新模型参数
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0

print('Finished Training')