基于 PyTorch 的 CNN 图片分类学习

在学习深度学习的过程中，图片分类是一个比较典型的任务。它的基本过程就是将图片输入神经网络，让模型学习图片中的特征，最后判断图片属于哪个类别。本次使用 PyTorch 搭建一个简单的 CNN 图片分类模型，通过食品图片数据集学习从数据读取到模型训练、测试的完整过程。

一、数据集的准备与读取

在训练模型之前，首先需要整理图片数据。代码通过 os.walk() 遍历 food_dataset2 文件夹中的图片，并根据图片所在的文件夹自动生成标签，最后生成 train.txt 和 test.txt。

例如：

复制代码

C:\...\train\apple\1.jpg 0
C:\...\train\banana\1.jpg 1

每一行前面是图片路径，后面是图片对应的类别编号。这样在后面读取数据时，只需要读取文本文件，就可以知道图片的位置和对应标签。

在 PyTorch 中，可以通过继承 Dataset 来创建自己的数据集：

复制代码

class food_dataset(Dataset):

这里最重要的是 __len__() 和 __getitem__() 两个方法。__len__() 用来返回数据集的大小，而 __getitem__() 根据索引读取指定的图片和标签。

例如：

复制代码

def __len__(self):
    return len(self.imgs)

表示数据集中有多少张图片。

复制代码

def __getitem__(self, idx):
    image = Image.open(self.imgs[idx])
    ...
    return image, label

表示获取某一张图片和它对应的标签。

代码前面的 USE_getitem 类也是为了理解这两个特殊方法。使用 p[1] 时，Python 会自动调用 __getitem__(1)；使用 len(p) 时，会自动调用 __len__()。PyTorch 的 Dataset 正是利用这种机制实现数据的索引读取。

读取图片后，还需要对图片进行预处理。代码使用：

复制代码

transforms.Resize([256,256])
transforms.ToTensor()

将不同大小的图片统一调整成 256×256，再转换成 Tensor。RGB 图片进入网络后，其数据形式可以理解为 3×256×256，其中3代表RGB三个颜色通道。

最后使用 DataLoader：

复制代码

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

将数据按照每批64张图片的方式提供给模型。shuffle=True 表示每轮训练时将数据打乱，有利于模型进行训练。

二、CNN卷积神经网络

数据准备完成后，就可以建立 CNN 模型。CNN 的主要作用是从图片中提取特征。代码中的网络主要由卷积层、ReLU激活函数、池化层和全连接层组成。

第一层：

复制代码

nn.Conv2d(3, 16, 5, 2, 2)

输入是3个通道的RGB图片，使用16个卷积核提取特征。之后使用 ReLU 激活函数增加网络的非线性能力，再通过最大池化：

复制代码

nn.MaxPool2d(2)

缩小特征图。

后面的 conv2 和 conv3 继续进行特征提取。随着网络不断加深，特征通道数量从：

复制代码

3 → 16 → 32 → 128

逐渐增加，而图片的空间尺寸逐渐减小。可以简单理解为，前面的网络学习图片中的一些基础特征，后面的网络逐渐学习更加复杂的特征。

经过卷积层后，需要将特征图转换成一维数据，再交给全连接层：

复制代码

x = x.view(x.size(0), -1)

最后通过：

复制代码

self.out = nn.Linear(..., 20)

输出20个分类结果。这里的20表示数据集一共有20个类别，模型会从这20个类别中选择一个作为最终预测结果。

需要注意的是，按照当前代码中的卷积和池化计算，conv3 后的尺寸是 128×64×64，因此全连接层的输入应该与实际尺寸对应。如果保持当前网络结构，应该使用：

复制代码

self.out = nn.Linear(128*64*64, 20)

如果希望使用原来的 128*32*32，则需要在 conv3 后再增加一次 MaxPool2d(2)。

三、模型训练

模型建立完成后，需要定义损失函数和优化器：

复制代码

loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=0.001
)

图片分类任务使用 CrossEntropyLoss 比较合适，它用于计算模型预测结果和真实标签之间的差距。优化器使用 Adam，它根据计算得到的梯度不断调整模型参数。

训练过程是整个深度学习中非常重要的部分：

复制代码

for X, y in train_dataloader:

    output = model(X)

    loss = loss_fn(output, y)

    optimizer.zero_grad()

    loss.backward()

    optimizer.step()

首先将图片 X 输入模型进行前向传播，得到预测结果 output；然后计算 Loss；接着使用 zero_grad() 清除之前的梯度，再通过 loss.backward() 进行反向传播，计算模型参数的梯度；最后使用 optimizer.step() 更新参数。

这个过程会不断重复：

复制代码

输入图片
↓
模型预测
↓
计算Loss
↓
反向传播
↓
更新参数

代码中：

复制代码

for epoch in range(5):

表示整个训练集重复学习5轮。每轮结束后计算平均 Loss，并观察模型的训练情况。通常随着训练进行，Loss 会逐渐下降，说明模型正在不断学习图片中的特征。

四、模型测试与准确率

训练完成之后，需要使用测试集检查模型的实际分类效果。

首先：

复制代码

model.eval()

将模型切换到测试模式。

然后：

复制代码

with torch.no_grad():

关闭梯度计算，因为测试阶段只需要进行预测，不需要更新模型参数。

模型输出20个类别的结果：

复制代码

output = model(X)

再使用：

复制代码

pred = output.argmax(dim=1)

找到输出结果中最大的类别作为预测结果。

例如模型输出：

复制代码

[0.2, 0.5, 3.7, 1.2]

最大值是 3.7，对应类别2，那么模型就预测这张图片属于第2类。

最后通过：

复制代码

correct += (pred == y).sum().item()

统计预测正确的数量，并使用：

复制代码

accuracy = correct / total * 100

计算测试集准确率。

准确率越高，说明模型在测试集上的分类效果越好。

五、总结

通过这次 CNN 图片分类实验，可以将之前学习的 PyTorch 知识串联起来。首先使用 os.walk() 整理图片并生成标签文件，然后通过 Dataset 和 DataLoader 完成数据读取和批量加载，再使用 transforms 对图片进行预处理。

模型部分使用 CNN 提取图片特征，通过卷积层、ReLU、池化层逐步处理图片，最后通过全连接层完成分类。训练时使用交叉熵损失函数计算误差，Adam 优化器根据反向传播得到的梯度更新模型参数。训练完成后，再通过测试集计算准确率，判断模型的分类能力。

整个过程可以概括为：

复制代码

准备图片数据
↓
Dataset读取数据
↓
DataLoader批量加载
↓
CNN提取图片特征
↓
计算Loss
↓
反向传播
↓
更新参数
↓
测试模型
↓
计算准确率

这也是使用 PyTorch 完成图像分类任务的一个基本流程，为后续学习数据增强、GPU训练、模型保存以及单张图片预测等内容打下基础。