PyTorch CNN图片分类与数据增强实践

一、图片数据增强

图片分类中，如果训练图片数量比较少，模型容易出现过拟合，也就是模型在训练集上表现很好，但是遇到没有见过的测试图片时准确率下降。为了提高模型的泛化能力，本次在训练数据中加入了数据增强。

训练集使用：

复制代码

data_transforms = {
    'trainda': transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.RandomRotation(45),
        transforms.CenterCrop(256),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.5),
        transforms.ColorJitter(
            brightness=0.2,
            contrast=0.1,
            saturation=0.1,
            hue=0.1
        ),
        transforms.RandomGrayscale(p=0.1),
        transforms.ToTensor(),
        transforms.Normalize(
            [0.485, 0.456, 0.406],
            [0.229, 0.224, 0.225]
        )
    ])
}

首先通过 Resize 将图片统一调整为 256×256，然后使用 RandomRotation 随机旋转图片，RandomHorizontalFlip 和 RandomVerticalFlip 分别进行水平和垂直翻转。

ColorJitter 可以随机改变图片的亮度、对比度、饱和度和色调，使同一张图片产生不同的变化。RandomGrayscale 则可以让一部分图片随机变成灰度图。

这些操作不会真正增加新的图片文件，而是在模型训练过程中随机改变图片，相当于让模型看到更多不同形式的训练样本，从而减少对某些固定图片特征的依赖。

训练集最后使用：

复制代码

transforms.ToTensor()

将图片转换成 Tensor，再通过 Normalize 对图片进行标准化。

需要注意的是，训练集使用了随机数据增强，而测试集没有使用随机旋转、翻转等操作：

复制代码

'valid': transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(
        [0.485, 0.456, 0.406],
        [0.229, 0.224, 0.225]
    )
])

测试集保持相对稳定，这样得到的测试准确率才更加具有参考意义。

二、Dataset、DataLoader与CNN模型

数据增强设置完成后，仍然使用自定义的 Dataset 读取图片。

复制代码

class food_dataset(Dataset):

通过 __len__() 获取数据集大小，通过 __getitem__() 根据索引读取图片和标签。

这里：

复制代码

image = Image.open(
    self.imgs[idx]
).convert('RGB')

将图片读取后统一转换成 RGB 格式，避免部分图片是灰度图或其他颜色模式导致输入通道不一致。

之后使用：

复制代码

train_dataloader = DataLoader(
    training_data,
    batch_size=64,
    shuffle=True
)

每次读取64张图片进行训练。训练数据使用 shuffle=True 打乱顺序，而测试数据使用 shuffle=False，保证测试过程更加稳定。

CNN模型仍然由卷积层、ReLU激活函数、池化层和全连接层组成。输入图片大小为 3×256×256，经过多次卷积和池化后提取图片特征，最后通过全连接层输出20个分类结果。

模型结构可以简单理解为：

复制代码

256×256 RGB图片
       ↓
卷积 + ReLU + 池化
       ↓
卷积 + ReLU + 池化
       ↓
卷积 + ReLU
       ↓
展平
       ↓
全连接层
       ↓
20个分类结果

这里的20表示数据集有20个类别。

三、模型训练与测试

本次代码增加了设备自动选择：

复制代码

device = 'cuda' if torch.cuda.is_available() else \
         'mps' if torch.backends.mps.is_available() else \
         'cpu'

程序会优先使用 CUDA GPU，如果没有 CUDA，则尝试使用 MPS，最后使用 CPU。

模型也需要移动到对应设备：

复制代码

model = CNN().to(device)

同时训练时：

复制代码

x, y = x.to(device), y.to(device)

将图片和标签放到相同的设备上。

训练过程被单独写成了 train() 函数。每一个 batch 都需要经过前向传播、计算损失、清除梯度、反向传播和更新参数：

复制代码

前向传播
↓
计算Loss
↓
梯度清零
↓
反向传播
↓
更新参数

代码同时统计了训练集的准确率：

复制代码

correct += (
    pred.argmax(1) == y
).type(torch.float).sum().item()

最后计算：

复制代码

accuracy = correct / size

测试过程与训练过程类似，但是测试时使用：

复制代码

model.eval()

with torch.no_grad():

其中 eval() 表示进入测试模式，no_grad() 表示测试时不计算梯度，因为测试阶段不需要更新模型参数。

本次设置：

复制代码

epochs = 20

表示模型训练20轮。每一轮训练结束后都会进行一次测试，并输出训练集和测试集的准确率以及 Loss。

四、保存最高准确率模型

本次代码相比之前比较重要的改进，是不再只保存最后一轮训练得到的模型，而是保存测试集准确率最高的模型。

首先：

复制代码

best_acc = 0

记录目前最高的准确率。

每轮测试完成后：

复制代码

if acc > best_acc:

判断当前准确率是否超过之前的最高准确率。

如果超过，就更新：

复制代码

best_acc = acc

然后保存模型参数：

复制代码

torch.save(
    model.state_dict(),
    './best2026-910.pth'
)

这样即使后面的训练出现准确率下降，也不会影响之前保存的最佳模型。

代码还进一步使用 TorchScript 保存模型：

复制代码

best_model = CNN().to('cpu')
best_model.load_state_dict(model.state_dict())
best_model.eval()

example_input = torch.randn(
    1, 3, 256, 256
)

script_model = torch.jit.trace(
    best_model,
    example_input
)

torch.jit.save(
    script_model,
    './best910.pth'
)

最终生成两个模型文件：

复制代码

best2026-910.pth
best910.pth

其中 best2026-910.pth 保存的是模型参数，而 best910.pth 是通过 TorchScript 保存的模型，可以用于后续模型部署和预测。

五、总结

本次实验是在基础 CNN 图片分类的基础上进行改进。通过 RandomRotation、随机翻转、ColorJitter、随机灰度等方法进行数据增强，使训练图片具有更多变化，减少模型过拟合的可能。

同时使用 Normalize 对图片进行标准化，并通过 CUDA 等设备加快模型训练。训练过程中将训练和测试分别封装成函数，每轮训练后计算准确率和 Loss，并通过 best_acc 自动保存测试集表现最好的模型。

整个程序的主要流程可以总结为：

复制代码

图片数据
↓
数据增强与标准化
↓
Dataset
↓
DataLoader
↓
CNN特征提取
↓
计算Loss
↓
反向传播
↓
更新参数
↓
测试准确率
↓
保存最高准确率模型

通过这次实验，不仅学习了 CNN 图片分类的基本方法，也进一步掌握了数据增强、GPU训练、模型评估和模型保存等深度学习中比较常用的技术。