深度学习之第六课卷积神经网络 (CNN)如何保存和使用最优模型

简介

        上一篇博客我们使用自己的数据集进行训练，生成了自己数据集的train.txt和test.txt文本,也说明了什么是数据增强。今天我们就来看看如何保存和使用最优模型。

深度学习之第五课卷积神经网络 (CNN)如何训练自己的数据集（食物分类）

一、保存最优模型

        对于训练结果，我们使用的都是的都是最后一轮的结果，但是最后一轮并不是训练出最好的模型，使用我们要把这个最好的模型保存下来。

1. 导入必要的库

import torch
from torch.utils.data import DataLoader,Dataset
from PIL import Image
from torchvision import transforms  # 对数据进行处理
import numpy as np
from torch import nn

• torch: PyTorch 深度学习框架的核心库
• DataLoader, Dataset: 用于数据加载和处理的工具类
• Image: 用于图像读取的 PIL 库
• transforms: 用于图像预处理的工具
• numpy: 用于数值计算的库
• nn: PyTorch 的神经网络模块

2. 图像预处理定义

data_transforms={     # 字典，存储不同阶段的图像处理方式
    'train':
        transforms.Compose([ # 组合多个变换
        transforms.Resize([300,300]),# 图像变换大小
        transforms.RandomRotation(45),# 图片随机旋转(-45到45度)
        transforms.CenterCrop(256),# 从中心开始裁剪到256x256
        transforms.RandomHorizontalFlip(p=0.5),# 50%概率水平翻转
        transforms.RandomVerticalFlip(p=0.5),# 50%概率垂直翻转
        transforms.ToTensor(), # 转换为Tensor格式
        transforms.Normalize([0.485,0.456,0.486],[0.229,0.224,0.225])# 归一化
        ]),
    'valid':
        transforms.Compose([
        transforms.Resize([256, 256]),
        transforms.ToTensor(),
        transforms.Normalize([0.485,0.456,0.486],[0.229,0.224,0.225])
        ]),
}

        定义了训练集和验证集的图像预处理流程，训练集使用了更多的数据增强技术来提高模型的泛化能力。

3. 自定义数据集类

class food_dataset(Dataset):# 继承Dataset类
    def __init__(self,file_path,transform=None):
        self.file_path=file_path
        self.imgs=[]
        self.labels=[]
        self.transform=transform
        # 读取文件列表和标签
        with open(file_path, 'r') as f:
            samples=[x.strip().split(' ') for x in f.readlines()]
            for img_path,label in samples:
                self.imgs.append(img_path)
                self.labels.append(label)

    def __len__(self):# 返回数据集大小
        return len(self.imgs)

    def __getitem__(self, idx):
        # 读取图像
        image=Image.open(self.imgs[idx])
        if self.transform: # 应用预处理
            image=self.transform(image)

        # 处理标签
        label=self.labels[idx]
        label=torch.from_numpy(np.array(label,dtype=np.int64))
        return image,label

自定义数据集类，用于加载食品图像数据和对应的标签，实现了 PyTorch 要求的三个核心方法：

• __init__: 初始化数据集，读取图像路径和标签
• __len__: 返回数据集样本数量
• __getitem__: 根据索引返回单个样本（图像和标签）

4. 数据加载器

# 创建数据集实例
train_data=food_dataset(file_path='train.txt',transform=data_transforms['train'])
test_data=food_dataset(file_path='test.txt',transform=data_transforms['train'])

# 创建数据加载器
train_dataloader=DataLoader(train_data,batch_size=32,shuffle=True)
test_dataloader=DataLoader(test_data,batch_size=32,shuffle=True)

• 使用自定义的food_dataset类创建训练集和测试集
• DataLoader用于批量加载数据，支持自动批处理、打乱数据等功能

5. 设备配置

device="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using{device},device")

自动选择可用的计算设备：

• 优先使用 NVIDIA GPU (cuda)
• 其次使用 Apple M 系列芯片的 GPU (mps)
• 最后使用 CPU

6. 定义 CNN 模型

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一个卷积块
        self.conv1=nn.Sequential(
            nn.Conv2d(in_channels=3,  # 输入通道数(RGB图像)
                      out_channels=32, # 输出通道数
                      kernel_size=5,   # 卷积核大小
                      stride=1,        # 步长
                      padding=2),      # 填充
            nn.ReLU(),               # 激活函数
            nn.MaxPool2d(2)          # 最大池化
        )
        # 第二个卷积块
        self.conv2=nn.Sequential(
            nn.Conv2d(32,64,5,1,2),
            nn.ReLU(),
            nn.Conv2d(64,64,5,1,2),
            nn.MaxPool2d(2)
        )
        # 第三个卷积块
        self.conv3=nn.Sequential(
            nn.Conv2d(64,128,5,1,2),
            nn.ReLU()
        )

        # 全连接层，输出20类
        self.out = nn.Linear(128 * 64 * 64, 20)

    def forward(self, x):
        # 前向传播
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        # 展平特征图
        x=x.view(x.size(0),-1)
        output=self.out(x)
        return output
model=CNN().to(device)

定义了一个简单的卷积神经网络模型，包含 3 个卷积块和 1 个全连接层，用于 20 类食品的分类任务。

7. 训练函数

def train(dataloader,model,loss_fn,optimizer):
    model.train()# 切换到训练模式
    batch_size_num=1 
    for X,y in dataloader:
        X,y=X.to(device),y.to(device)  # 将数据移到计算设备
        pred=model.forward(X)          # 前向传播，得到预测结果
        loss=loss_fn(pred,y)           # 计算损失

        optimizer.zero_grad()          # 梯度清零
        loss.backward()                # 反向传播计算梯度
        optimizer.step()               # 更新参数

        loss=loss.item()
        # 每64个批次打印一次损失
        if batch_size_num % 64 == 0:
            print(f"loss:{loss:>7f} [number:{batch_size_num}]")
        batch_size_num+=1

训练函数实现了模型训练的基本流程：

• 前向传播计算预测结果
• 计算损失
• 反向传播计算梯度
• 更新模型参数

8. 测试函数

best_acc=0
def test(dataloader,model,loss_fn):
    size=len(dataloader.dataset)
    num_batches=len(dataloader)
    model.eval()  # 切换到评估模式
    test_loss,correct=0,0
    with torch.no_grad():  # 禁用梯度计算，节省内存
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss+=loss_fn(pred,y).item()
            # 计算正确预测的数量
            correct+=(pred.argmax(1)==y).type(torch.float).sum().item()
    
    # 计算平均损失和准确率
    test_loss/=num_batches
    correct/=size
    print(f"Test resilt:\n Accuracy:{(100*correct)}%,Avg loss:{test_loss}")

    # 保存最优模型
    global best_acc
    if correct > best_acc :
        best_acc=correct
        torch.save(model,'best.pt')  # 保存整个模型

保存最优模型

        保存最优模型是通过测试函数进行保存，有两种方式，一种是保存模型的参数（不包含模型的结构）内存比较小。另外一种就是保存整个模型包括模型结构与参数。内存就相对于大点。常见的格式有.pt和.pth

#保存最优模型
    global best_acc
    if correct > best_acc :
        best_acc=correct
        # torch.save(model.state_dict(),'best.pth') #保存参数
        torch.save(model,'best.pt') # 保存整个模型

9.训练配置和执行

# 定义损失函数和优化器
loss_fn=nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(model.parameters(),lr=0.001)
# 学习率调度器，每10个epoch学习率减半
scheduler=torch.optim.lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.5)

# 训练轮次
epochs=150
acc_s=[]
loss_s=[]
for t in range(epochs):
    print(f"Epoch{t+1}\n...............")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
    scheduler.step()# 调整学习率
print("Done!")

10.绘制训练曲线

from matplotlib import pyplot as plt

# 创建一个1行2列的子图布局，返回图形对象和子图数组
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))  # 可以指定figsize调整大小

# 第一个子图：准确率曲线
ax1.plot(range(0, epochs), acc_s)
ax1.set_xlabel('epoch')
ax1.set_ylabel('accuracy')
ax1.set_title('Accuracy Curve')  # 可以添加标题

# 第二个子图：损失曲线
ax2.plot(range(0, epochs), loss_s)
ax2.set_xlabel('epoch')
ax2.set_ylabel('loss')
ax2.set_title('Loss Curve')  # 可以添加标题

plt.tight_layout()  # 自动调整子图间距
plt.savefig('training_curves.png', dpi=300, bbox_inches='tight')
plt.show()
print("done!")

二、使用最优模型

1.基础库导入

import torch
from torch.utils.data import DataLoader, Dataset
from PIL import Image
from torchvision import transforms
import numpy as np
from torch import nn

2.测试集图像预处理配置

data_transforms = {
    'valid': transforms.Compose([
        transforms.Resize([256, 256]),  # 缩放图像到256×256像素
        transforms.ToTensor(),          # 转换为PyTorch张量（格式：[C, H, W]，数值归一化到0-1）
        transforms.Normalize([0.485, 0.456, 0.486], [0.229, 0.224, 0.225])  # 标准化
    ])
}

• 键'valid'对应测试集的预处理流程（与训练时测试集的处理逻辑完全一致，避免数据分布差异影响预测结果）；
• transforms.Compose：将多个预处理操作 “串联” 成一个流水线；
• 关键操作说明：
  1. Resize([256,256])：将所有测试图像统一缩放到 256×256，确保输入模型的尺寸一致；
  2. ToTensor()：将 PIL 图像（H×W×C，像素值 0-255）转换为 PyTorch 张量（C×H×W，像素值归一化到 0-1），符合模型输入格式；
  3. Normalize：使用 ImageNet 数据集的均值和标准差对张量标准化（训练时模型以此为基准，测试时需保持一致），加速模型推理并提升稳定性。

3.自定义测试集数据集类

class food_dataset(Dataset):
    def __init__(self, file_path, transform=None):
        self.file_path = file_path  # 测试集文件列表的txt路径
        self.imgs = []              # 存储所有测试图像的路径
        self.labels = []            # 存储所有测试图像的真实标签
        self.transform = transform  # 图像预处理流水线
        # 读取txt文件，解析图像路径和真实标签
        with open(file_path, 'r') as f:
            # 按行读取txt，每行格式为“图像路径 真实标签”，分割后存入samples
            samples = [x.strip().split(' ') for x in f.readlines()]
            for img_path, label in samples:
                self.imgs.append(img_path)  # 收集图像路径
                self.labels.append(label)   # 收集真实标签

    def __len__(self):
        # 返回测试集的样本总数（必须实现的Dataset抽象方法）
        return len(self.imgs)

    def __getitem__(self, idx):
        # 根据索引idx获取单个测试样本（图像+真实标签，必须实现的Dataset抽象方法）
        # 1. 读取图像
        image = Image.open(self.imgs[idx])  # 用PIL打开指定路径的图像
        # 2. 应用预处理
        if self.transform:
            image = self.transform(image)   # 将图像按预处理流水线处理为张量
        # 3. 处理真实标签：转换为int64类型的PyTorch张量（符合模型标签格式）
        label = torch.from_numpy(np.array(self.labels[idx], dtype=np.int64))
        # 返回处理后的“图像张量”和“真实标签张量”
        return image, label

4.CNN 模型结构定义

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()  # 继承nn.Module的初始化方法
        # 第1个卷积块：卷积+ReLU激活+最大池化（提取低级特征，如边缘、纹理）
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 16, 5, 1, 2),  # 卷积层：输入3通道（RGB），输出16通道，卷积核5×5，步长1，填充2
            nn.ReLU(),                  # ReLU激活函数：引入非线性，增强模型表达能力
            nn.MaxPool2d(2)             # 最大池化：2×2池化核，将特征图尺寸缩小1/2（保留关键特征，减少计算量）
        )
        # 第2个卷积块：2次卷积+ReLU激活+最大池化（提取中级特征，如局部形状）
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2), # 卷积层：输入16通道，输出32通道
            nn.ReLU(),
            nn.Conv2d(32, 32, 5, 1, 2), # 卷积层：输入32通道，输出32通道（加深特征提取）
            nn.MaxPool2d(2)             # 再次池化，特征图尺寸再缩小1/2
        )
        # 第3个卷积块：卷积+ReLU激活（提取高级特征，如物体部件）
        self.conv3 = nn.Sequential(
            nn.Conv2d(32, 64, 5, 1, 2), # 卷积层：输入32通道，输出64通道
            nn.ReLU()
        )
        # 全连接层：将卷积提取的特征图“展平”后，映射到20个类别（食品分类共20类）
        self.out = nn.Linear(64 * 64 * 64, 20)  # 输入维度=64（通道数）×64×64（特征图尺寸）

    def forward(self, x):
        # 前向传播：定义数据在模型中的流动路径（与训练时完全一致）
        x = self.conv1(x)  # 经过第1个卷积块
        x = self.conv2(x)  # 经过第2个卷积块
        x = self.conv3(x)  # 经过第3个卷积块
        x = x.view(x.size(0), -1)  # 特征图展平：(batch_size, 64*64*64)，-1表示自动计算剩余维度
        output = self.out(x)  # 经过全连接层，输出每个类别的预测分数
        return output

5.计算设备配置

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

6.加载已训练好的模型

model = torch.load('best.pt').to(device)
model.eval()

• torch.load('best.pt')：加载训练时保存的 “最优模型”（best.pt是训练过程中准确率最高的模型文件）；
• .to(device)：将加载的模型移动到指定计算设备（与测试数据设备一致，避免数据 / 模型设备不匹配报错）；
• model.eval()：将模型切换为评估模式（关键操作）：
  • 禁用 Dropout 层（本模型无 Dropout，但需保持规范）；
  • 固定 BatchNorm 层的统计参数（避免测试时更新均值 / 方差，影响预测结果）。

7.加载测试集数据

test_data = food_dataset(file_path='test.txt', transform=data_transforms['valid'])
test_dataloader = DataLoader(test_data, batch_size=32, shuffle=True)

8.测试函数与结果计算

result = []  # 存储所有测试样本的预测标签
labels = []  # 存储所有测试样本的真实标签
def test_ture(dataloader, model):
    with torch.no_grad():  # 禁用梯度计算（关键！测试时无需反向传播，节省内存并加速）
        for X, y in dataloader:  # 遍历测试集的每个批次
            X, y = X.to(device), y.to(device)  # 将批次数据移动到指定设备
            pred = model(X)  # 模型推理：输入图像张量，输出每个类别的预测分数
            result.extend(pred.argmax(1).tolist())  # 取预测分数最大的类别作为预测标签，存入result
            labels.extend(y.tolist())  # 将真实标签转换为列表，存入labels

    # 计算准确率：（预测正确的样本数 / 总样本数）× 100%
    correct = sum(p == l for p, l in zip(result, labels))  # 统计预测标签与真实标签一致的数量
    total = len(labels)  # 测试集总样本数
    accuracy = correct / total * 100 if total > 0 else 0  # 避免除以0（空测试集时准确率为0）

    return accuracy

9.执行测试与输出结果

# 执行测试并获取准确率
accuracy = test_ture(test_dataloader, model)

# 输出结果
print('预测值：\t', result)
print('真实值：\t', labels)
print(f'准确率：\t{accuracy:.2f}%')