深度学习之第六课卷积神经网络 (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 nntorch: 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),# 从中心开始裁剪到256x256transforms.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_pathself.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,0with 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_batchescorrect/=sizeprint(f"Test resilt:\n Accuracy:{(100*correct)}%,Avg loss:{test_loss}")# 保存最优模型global best_accif correct > best_acc :best_acc=correcttorch.save(model,'best.pt') # 保存整个模型保存最优模型
保存最优模型是通过测试函数进行保存，有两种方式，一种是保存模型的参数（不包含模型的结构）内存比较小。另外一种就是保存整个模型包括模型结构与参数。内存就相对于大点。常见的格式有.pt和.pth
#保存最优模型global best_accif 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 nn2.测试集图像预处理配置
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：将多个预处理操作 “串联” 成一个流水线；- 关键操作说明：
Resize([256,256])：将所有测试图像统一缩放到 256×256，确保输入模型的尺寸一致；ToTensor()：将 PIL 图像（H×W×C，像素值 0-255）转换为 PyTorch 张量（C×H×W，像素值归一化到 0-1），符合模型输入格式；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，每行格式为“图像路径 真实标签”，分割后存入samplessamples = [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, label4.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，填充2nn.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 output5.计算设备配置
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()) # 取预测分数最大的类别作为预测标签，存入resultlabels.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 accuracy9.执行测试与输出结果
# 执行测试并获取准确率
accuracy = test_ture(test_dataloader, model)# 输出结果
print('预测值：\t', result)
print('真实值：\t', labels)
print(f'准确率：\t{accuracy:.2f}%')