2025-05-31 Python深度学习10——模型训练流程

本文环境：

• Pycharm 2025.1
• Python 3.12.9
• Pytorch 2.6.0+cu124

​ 本文以 CIFAR-10 为例，介绍模型的大致训练流程。相关的 Python 包如下：

import torch
import torchvision
from torch import nn, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time

1 数据准备

1.1 下载与预处理

​ 使用torchvision.datasets.CIFAR10下载 CIFAR-10 数据集（32x32 彩色图像，10 类），分为训练集（train=True，5 万张）和测试集（train=False，1 万张）。

# 准备数据集
train_data = torchvision.datasets.CIFAR10(root='./dataset',train=True,transform=torchvision.transforms.ToTensor(),download=True
)test_data = torchvision.datasets.CIFAR10(root='./dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True
)

• transform=torchvision.transforms.ToTensor()：将图像转为 PyTorch 张量（Tensor），并自动归一化到 [0, 1] 范围。
• download=True：若本地无数据，自动下载。

1.2 数据加载

​ 通过DataLoader分批次加载数据：

# 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

• batch_size=64：每批次处理 64 张图片，平衡内存占用和训练效率。
• 训练集默认不打乱（未设置shuffle），测试集可添加shuffle=True以增强评估可靠性。

2 模型构建

2.1 自定义 CNN 模型

MyModel是一个3层卷积神经网络（CNN）：

1. 卷积层：nn.Conv2d(3, 32, 5, 1, 2)

   输入通道 3（RGB），输出通道 32，5×5 卷积核，步长 1，填充 2（保持尺寸不变）。

2. 池化层：nn.MaxPool2d(2)

   2×2最大池化，尺寸减半。

3. 全连接层

   • nn.Linear(64 * 4 * 4, 64)

     将展平后的特征（64 通道×4×4尺寸）映射到 64 维。

   • nn.Linear(64, 10)

     最终输出 10 类。

class MyModel(nn.Module):def __init__(self):super(MyModel, self).__init__()self.model = nn.Sequential(nn.Conv2d(3, 32, 5, 1, 2),nn.MaxPool2d(2),nn.Conv2d(32, 32, 5, 1, 2),nn.MaxPool2d(2),nn.Conv2d(32, 64, 5, 1, 2),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64 * 4 * 4, 64),nn.Linear(64, 10))def forward(self, x):return self.model(x)

2.2 GPU加速

​ 通过.to(device)将模型和数据移至 GPU（若可用），显著加速计算。

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = MyModel().to(device)  # 使用GPU

3 训练配置

3.1 损失函数

​ 使用nn.CrossEntropyLoss()，适用于多分类任务，计算预测概率与真实标签的交叉熵。

# 损失函数
loss_fn = nn.CrossEntropyLoss().to(device)  # 使用GPU

3.2 优化器

​ 使用torch.optim.SGD，随机梯度下降，学习率lr=1e-2，控制参数更新步长。

# 损失函数
loss_fn = nn.CrossEntropyLoss().to(device)  # 使用GPU

3.3 训练参数

1. total_train_step：记录训练次数，用于日志和调试。
2. total_test_step：记录测试次数，用于日志和调试。
3. epoch=20：遍历完整数据集 20 次。

# 设置训练网络的一些参数
total_train_step = 0  # 记录训练次数
total_test_step = 0  # 记录测试次数
epoch = 20  # 训练的轮数

4 训练循环

数据加载  →  模型初始化  →  训练循环  →  测试评估  →  保存模型↑          ↑                  ↓           ↓└───TensorBoard日志 ←────── 参数更新 ←── 梯度计算

4.1 训练模式 (model.train())

1. 前向传播：输入图像imgs，模型输出预测outputs。
2. 计算损失：loss = loss_fn(outputs, targets)，衡量预测误差。
3. 反向传播：
   • optimizer.zero_grad()：清空梯度，避免累积。
   • loss.backward()：计算梯度（链式法则）。
   • optimizer.step()：更新模型参数。
4. 日志记录：每 100 次训练记录损失和时间到 TensorBoard 中。

for i in range(epoch):print(f"------------第 {i + 1} 轮训练开始------------")# 训练步骤开始model.train()for data in train_dataloader:imgs, targets = dataimgs = imgs.to(device)  # 使用GPUtargets = targets.to(device)  # 使用GPUoutputs = model(imgs)loss = loss_fn(outputs, targets)# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()total_train_step += 1if total_train_step % 100 == 0:end_time = time.time()print(f"第 {total_train_step} 次训练，Loss：{loss.item()}，Time：{end_time - start_time}")writer.add_scalar("train_loss", loss.item(), total_train_step)start_time = time.time()

4.2 评估模式 (model.eval())

1. 关闭梯度计算：with torch.no_grad()，节省内存并加速。
2. 测试指标：
   • 总损失：累加所有批次的损失total_test_loss。
   • 准确率：统计预测正确的样本数（outputs.argmax(dim=1) == targets）。
3. 日志记录：每轮测试后保存损失和准确率到 TensorBoard。
4. 保存模型：通过torch.save()方法将模型的 state_dict 保存到本地文件中。

    # 测试步骤开始model.eval()total_test_loss = 0total_accuracy = 0accuracy_rate = 0with torch.no_grad():for data in test_dataloader:imgs, targets = dataimgs = imgs.to(device)  # 使用GPUtargets = targets.to(device)  # 使用GPUoutputs: Tensor = model(imgs)loss = loss_fn(outputs, targets)total_test_loss += lossaccuracy = outputs.argmax(dim=1) == targetstotal_accuracy += accuracy.sum()total_test_step += 1accuracy_rate = total_accuracy / test_data_sizeprint(f"第 {i + 1} 轮测试，Loss：{total_test_loss}，Accuracy：{total_accuracy} ({accuracy_rate})")writer.add_scalar("test_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_accuracy, total_test_step)writer.add_scalar("accuracy_rate", accuracy_rate, total_test_step)torch.save(model.state_dict(), f"model/my_model.pth")  # 保存模型writer.close()

说明

1. 训练与评估模式切换

   • model.train()：启用 Dropout 和 BatchNorm 的训练行为（如随机丢弃神经元）。

   • model.eval()：固定 Dropout 和 BatchNorm 的统计量，确保评估一致性。

2. GPU 数据迁移

   需将输入数据 imgs 和标签 targets 均移至 GPU，否则会报错。

3. 梯度清零

   避免梯度累加导致参数更新错误。

完整代码

# train_gpu_2.pyimport torch
import torchvision
from torch import nn, Tensor
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import time# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 准备数据集
train_data = torchvision.datasets.CIFAR10(root='./dataset',train=True,transform=torchvision.transforms.ToTensor(),download=True
)test_data = torchvision.datasets.CIFAR10(root='./dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True
)# 数据集大小
train_data_size = len(train_data)
test_data_size = len(test_data)print(f"训练集数量：{train_data_size}")
print(f"测试集数量：{test_data_size}")# 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 创建网络模型
class MyModel(nn.Module):def __init__(self):super(MyModel, self).__init__()self.model = nn.Sequential(nn.Conv2d(3, 32, 5, 1, 2),nn.MaxPool2d(2),nn.Conv2d(32, 32, 5, 1, 2),nn.MaxPool2d(2),nn.Conv2d(32, 64, 5, 1, 2),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64 * 4 * 4, 64),nn.Linear(64, 10))def forward(self, x):return self.model(x)model = MyModel().to(device)  # 使用GPU# 损失函数
loss_fn = nn.CrossEntropyLoss().to(device)  # 使用GPU# 优化器
lr = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr=lr)# 设置训练网络的一些参数
total_train_step = 0  # 记录训练次数
total_test_step = 0  # 记录测试次数
epoch = 20  # 训练的轮数# 添加 tensorboard
writer = SummaryWriter("../logs_train")for i in range(epoch):print(f"------------第 {i + 1} 轮训练开始------------")start_time = time.time()# 训练步骤开始model.train()for data in train_dataloader:imgs, targets = dataimgs = imgs.to(device)  # 使用GPUtargets = targets.to(device)  # 使用GPUoutputs = model(imgs)loss = loss_fn(outputs, targets)# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()total_train_step += 1if total_train_step % 100 == 0:end_time = time.time()print(f"第 {total_train_step} 次训练，Loss：{loss.item()}，Time：{end_time - start_time}")writer.add_scalar("train_loss", loss.item(), total_train_step)start_time = time.time()# 测试步骤开始model.eval()total_test_loss = 0total_accuracy = 0accuracy_rate = 0with torch.no_grad():for data in test_dataloader:imgs, targets = dataimgs = imgs.to(device)  # 使用GPUtargets = targets.to(device)  # 使用GPUoutputs: Tensor = model(imgs)loss = loss_fn(outputs, targets)total_test_loss += lossaccuracy = outputs.argmax(dim=1) == targetstotal_accuracy += accuracy.sum()total_test_step += 1accuracy_rate = total_accuracy / test_data_sizeprint(f"第 {i + 1} 轮测试，Loss：{total_test_loss}，Accuracy：{total_accuracy} ({accuracy_rate})")writer.add_scalar("test_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_accuracy, total_test_step)writer.add_scalar("accuracy_rate", accuracy_rate, total_test_step)torch.save(model.state_dict(), f"model/my_model.pth")  # 保存模型writer.close()

5 模型验证

​ 准备待验证的图片，放在 imgae 目录下。

​ 编写 test.py 文件，用于验证模型。

# test.pyimport torch
import torchvision
from PIL import Image
from torch import nn# 定义图片路径
image_path = "image/dog.png"# 打开图片并转换为RGB格式
image = Image.open(image_path).convert('RGB')# 定义图片转换操作
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),  # 将图片大小调整为32x32torchvision.transforms.ToTensor()  # 将图片转换为张量
])# 对图片进行转换操作
image = transform(image).reshape(1, 3, 32, 32)# 定义模型类
class MyModel(nn.Module):def __init__(self):super(MyModel, self).__init__()# 定义模型结构self.model = nn.Sequential(nn.Conv2d(3, 32, 5, 1, 2),  # 第一个卷积层，输入通道数为3，输出通道数为32，卷积核大小为5，步长为1，填充为2nn.MaxPool2d(2),  # 最大池化层，池化核大小为2nn.Conv2d(32, 32, 5, 1, 2),  # 第二个卷积层，输入通道数为32，输出通道数为32，卷积核大小为5，步长为1，填充为2nn.MaxPool2d(2),  # 最大池化层，池化核大小为2nn.Conv2d(32, 64, 5, 1, 2),  # 第三个卷积层，输入通道数为32，输出通道数为64，卷积核大小为5，步长为1，填充为2nn.MaxPool2d(2),  # 最大池化层，池化核大小为2nn.Flatten(),  # 展平操作nn.Linear(64 * 4 * 4, 64),  # 全连接层，输入维度为64*4*4，输出维度为64nn.Linear(64, 10)  # 全连接层，输入维度为64，输出维度为10)def forward(self, x):return self.model(x)# 加载模型参数
model_dict = torch.load('model/my_model.pth')
model = MyModel()
model.load_state_dict(model_dict)
model.to('cuda')# 设置模型为评估模式
model.eval()
# 关闭梯度计算
with torch.no_grad():# 将图片转换为GPU格式image = image.to('cuda')# 进行模型推理output = model(image)# 打印输出结果
print(output)
# 打印输出结果中最大值的索引
print(output.argmax(1))

​ 验证结果如下，表明 dog.png 图片的预测结果索引为 5，即第 6 类预测。

​ 依据分类规则，预测结果为 dog，是正确的。