用 PyTorch 做一个食品图片分类模型：从整理数据到保存模型

最近我学习了如何使用 PyTorch 完成图片分类，并尝试搭建了一个食品图片识别模型。这个项目的目标非常明确：给模型输入一张食品图片，让它判断这张图片属于哪一种食品。

听起来好像很复杂，但实际拆开以后，整个过程可以分成几个比较清楚的步骤：

整理图片
划分训练集和测试集
处理图片
创建数据集
搭建卷积神经网络
训练模型
测试准确率
保存效果最好的模型

这篇文章就按照这个顺序，记录一下整个食品分类模型是怎样完成的。

一、图片分类

图片分类的任务，就是让计算机判断一张图片属于哪个类别。

比如，现在的数据集中有二十种食品，那么模型看到一张图片后，需要输出二十个分数。每个分数对应一种食品，哪个类别的分数最高，模型就认为这张图片属于哪个类别。

假设模型输出的结果是：

[0.2, -0.5, 1.3, 4.8, 0.7, ...]

如果第4个位置的分数最大，那么模型就会把这张图片判断为第4类食品。

不过，模型一开始并不知道每种食品长什么样。我们需要先准备大量带有正确标签的图片，让模型反复学习。

二、整理食品图片

在正式训练之前，首先需要整理数据集。一般情况下，可以把不同种类的食品图片分别放在不同文件夹中。

目录结构可以类似这样：

每个文件夹代表一个食品类别。程序遍历目录时，会根据文件夹所在的位置生成标签。例如，第一个文件夹中的图片标签是0，第二个文件夹中的图片标签是1。

模型训练时不能只知道图片在哪里，还必须知道每张图片的正确类别。因此，代码会生成两个文本文件：

train.txt
test.txt

每一行保存一张图片的路径和对应标签，例如：

D:\food_dataset2\0\food_001.jpg 0
D:\food_dataset2\3\food_018.jpg 3
D:\food_dataset2\12\food_026.jpg 12

这里用数字代表具体标签

这样，程序读取一行内容，就可以同时知道图片路径和正确答案。

三、划分训练集和测试集

一份图片数据不能全部用来训练模型，还需要留出一部分图片检查模型效果。

要划分训练集和测试机

代码根据图片文件名最后的数字划分数据：

if (
    not os.path.splitext(file)[0].rsplit('_', 1)[-1].isdigit()
    or int(os.path.splitext(file)[0].rsplit('_', 1)[-1]) % 5 != 0
):
    train_txt.write(path_1 + ' ' + label + '\n')
else:
    test_txt.write(path_1 + ' ' + label + '\n')

简单来说，如果文件名最后的数字可以被5整除，就把图片放入测试集；其他图片放入训练集。

按照这种方法，大约80%的图片会进入训练集，剩下20%左右进入测试集，也就是常见的8:2划分。

四、对训练图片进行数据增强？

图片数量有限时，模型可能很快记住训练图片，却不一定能够识别新的图片。

例如，同一种食品在现实中可能出现很多变化：拍摄角度不一样；图片亮度不一样；食品在画面中的位置不一样；图片可能发生轻微旋转。

如果训练图片比较固定，模型可能只会识别特定角度下的食品。为了让模型理解更多变化，可以在训练时随机处理图片，这种方法叫作数据增强。

代码中的训练集处理方式如下：

data_transforms = {
    'traindata': transforms.Compose([
        transforms.RandomResizedCrop(
            256,
            scale=(0.8, 1.0),
            ratio=(0.9, 1.1)
        ),
        transforms.RandomRotation(15),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(
            brightness=0.2,
            contrast=0.2,
            saturation=0.2,
            hue=0.05
        ),

          # 食品图片不建议上下颠倒
          # transforms.RandomVerticalFlip(p=0.1),

        transforms.RandomGrayscale(p=0.05),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
}

这些操作分别有不同作用。

RandomResizedCrop 会随机裁剪图片，再调整为256×256。这样可以让食品在画面中的大小和位置出现一些变化。

RandomRotation(15) 会让图片在一定范围内随机旋转，最大角度约为15度。

RandomHorizontalFlip(p=0.5) 表示图片有50%的概率进行水平翻转。

ColorJitter 会对亮度、对比度、饱和度和色调进行少量随机调整。

RandomGrayscale(p=0.05) 表示图片有5%的概率变成灰度效果。

这些处理不会改变图片的实际类别，却可以让同一张图片在不同训练轮次中呈现出不同效果。这样相当于增加了训练数据的丰富程度。

食品图片通常不适合上下翻转，因为一盘食物倒过来可能与现实中的拍摄方式差别太大。、

这里我定义的是字典，traindata和valid是两个键，对应两个值，当然可以不用字典，直接定义两个变量也可以。

五、测试集不需要数据增强

测试集的处理方式比训练集简单：

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

测试图片只需要统一尺寸、转换成张量并完成标准化，不需要随机旋转、随机裁剪或改变颜色。

所以可以简单记成：

训练集：可以随机变化，让模型多练习
测试集：保持稳定，用来检查真实效果

六、ToTensor 和 Normalize 

普通图片不能直接送入 PyTorch 模型，需要先通过 ToTensor() 转换成张量。

转换之后，图片的结构通常是：

[通道数, 高度, 宽度]

对于彩色图片来说，形状一般为：

[3, 256, 256]

三个通道分别代表红色、绿色和蓝色。

图片原来的像素值通常在0到255之间，转换成张量后一般会变成0到1之间的小数。

接下来使用 Normalize 对图片进行标准化：

transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)

标准化可以把不同通道的数据调整到更适合模型训练的范围，让模型训练得更加稳定。这里是0~1之间。

这里采用的是均值和标准差。

七、自定义食品图片数据集

PyTorch 提供了 Dataset 类，我们可以在它的基础上创建自己的食品数据集。

class food_dataset(Dataset):
    def __init__(self, file_path, transform=None):
        self.file_path = file_path
        self.imgs = []
        self.labels = []
        self.transform = transform

        with open(self.file_path) as f:
            samples = [
                x.strip().rsplit(' ', 1)
                for x in f
                if x.strip()
            ]

        for img_path, label in samples:
            self.imgs.append(img_path)
            self.labels.append(label)

初始化时，程序打开 train.txt 或 test.txt，读取每一行中的图片路径和标签，然后分别保存到 self.imgs 和 self.labels 中。

自定义数据集还需要实现两个方法。

第一个是 __len__()：

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

它负责告诉程序数据集中一共有多少张图片。

第二个是 __getitem__()：

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

当程序需要第 idx 张图片时，这个方法会读取图片，执行对应的图片处理，再把标签转换成张量，最后返回图片和标签。

模型不需要关心图片具体放在哪里，只需要告诉数据集自己想要第几张图片，数据集就会把处理好的图片和正确标签交给模型。

八、使用 DataLoader 分批读取数据

创建数据集之后，还需要使用 DataLoader 分批读取：

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

test_dataloader = DataLoader(
    test_data,
    batch_size=32,
    shuffle=False
)

batch_size=32 表示模型每次处理32张图片。

如果一张一张训练，速度会比较慢；如果一次放入太多图片，又可能占用大量显存。因此，分批训练是在运行速度和显存占用之间进行平衡。

训练集设置了：

shuffle=True

表示每轮训练前会打乱图片顺序，防止模型总是按照固定顺序学习。

测试集设置了：

shuffle=False

因为测试时不需要打乱顺序，只需要稳定地计算最终结果。

九、搭建卷积神经网络

食品分类模型使用了卷积神经网络，也就是 CNN。

CNN 特别适合处理图片。它可以从图片中逐步提取特征：前面的卷积层可能学到边缘和颜色，后面的卷积层则可能学到食品的形状、纹理以及局部结构。

代码中的第一组卷积层为：

self.conv1 = nn.Sequential(
    nn.Conv2d(
        in_channels=3,
        out_channels=16,
        kernel_size=5,
        stride=1,
        padding=2
    ),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2)
)

输入图片有三个颜色通道，因此 in_channels=3。

out_channels=16 表示这一层会得到16张特征图。可以把每张特征图理解成模型从不同角度观察图片的结果：有的可能关注颜色，有的可能关注边缘，还有的可能关注纹理。

卷积完成后使用 ReLU 激活函数，为模型加入非线性能力。

接着使用最大池化，将特征图从256×256缩小到128×128。图片虽然变小了，但比较明显的特征仍然会被保留下来，同时计算量也会减少。

十、继续提取更深层的特征

第二组卷积层将通道数量从16增加到32：

self.conv2 = nn.Sequential(
    nn.Conv2d(16, 32, 5, 1, 2),
    nn.ReLU(),
    nn.Conv2d(32, 32, 5, 1, 2),
    nn.ReLU(),
    nn.MaxPool2d(2)
)

这一部分连续使用了两次卷积。前面的卷积层提取简单特征，后面的卷积层再把这些简单特征组合起来。

经过池化以后，特征图大小会从128×128缩小到64×64。

第三组卷积层继续把通道数量增加到128：

self.conv3 = nn.Sequential(
    nn.Conv2d(32, 128, 5, 1, 2),
    nn.ReLU()
)

到了这一层，模型处理的已经不是原始像素，而是前面卷积层提取出来的特征。它可能逐渐关注到食品的轮廓、表面纹理、颜色搭配以及局部结构。

十一、从特征图得到分类结果

卷积层输出的数据形状大约为：

[batch_size, 128, 64, 64]

但是，全连接层不能直接处理这种四维数据，因此需要先将其展开：

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

展开后，每张图片的所有特征会连接成一个长向量，再送入全连接层：

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

out_features=20 表示数据集中有20个食品类别。

最终，模型会为每张图片输出20个分数，然后选择分数最高的位置作为预测类别。

十二、CPU 和 GPU 的选择

代码会自动检查电脑是否支持 CUDA：

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

如果存在可用的 NVIDIA GPU，程序会使用 GPU 训练；如果没有，就使用 CPU。

训练过程中，模型、图片和标签必须放在同一个设备上：

model = CNN().to(device)
images = images.to(device)
labels = labels.to(device)

如果模型在 GPU 上，而图片还在 CPU 上，程序就会因为设备不同而报错。

十三、模型训练

训练函数首先让模型进入训练状态：

model.train()

然后从 DataLoader 中分批读取图片和标签：

for images, labels in dataloader:

模型根据图片得到预测结果：

pred = model(images)

接着通过损失函数比较预测结果和正确标签：

loss = loss_fn(pred, labels)

损失值可以理解成模型当前“错得有多严重”。损失越大，说明预测结果和正确答案差距越大；损失逐渐降低，通常说明模型正在学习。

更新模型参数时，需要执行三个步骤：

optimizer.zero_grad()
loss.backward()
optimizer.step()

zero_grad() 用来清除上一批数据留下的梯度。

backward() 根据当前错误计算每个参数应该向哪个方向调整。

step() 则真正完成参数更新。

整个过程会对训练集中的每一批图片执行一次，所有图片学习完一遍，就完成了一个 Epoch。

十四、测试模型效果

测试模型时，需要先切换到评估状态：

model.eval()

然后关闭梯度计算：

with torch.no_grad():

测试阶段只需要检查模型表现，不需要更新参数，所以关闭梯度可以减少内存占用，也能提高运行速度。

模型输出结果后，使用下面的代码找到分数最高的类别：

pred.argmax(1)

再将预测类别与正确标签比较：

correct += (
    pred.argmax(1) == labels
).sum().item()

最后，用预测正确的图片数量除以测试集图片总数，就能得到准确率：

accuracy = correct / size

例如测试集中有1000张图片，其中850张识别正确，那么准确率就是85%。

十五、保存效果最好的模型

模型每训练一轮，测试准确率都可能发生变化。有时准确率提高，有时也可能下降。

因此，代码使用 best_acc 记录历史最高准确率：

best_acc = 0.0

如果当前准确率超过之前的最好结果，就保存模型：

if accuracy > best_acc:
    best_acc = accuracy
    torch.save(
        model.state_dict(),
        weight_model_path
    )

这样，即使后面的训练结果变差，之前效果最好的模型也不会被覆盖。

代码同时保存了两种模型文件。

第一种只保存模型参数：

cnn_food_model.pth

这种文件通常体积较小，但加载时必须先创建相同的 CNN 网络结构。

第二种保存 TorchScript 模型：

cnn_food_model_scripted.pth

它包含可以直接执行的模型结构和参数，后续部署时会更加方便。

因为 Windows 中的中文路径有时会影响 TorchScript 保存，所以代码使用文件对象写入，尽量避开中文路径带来的问题。

十六、损失函数和优化器

模型使用交叉熵损失函数：

loss_function = nn.CrossEntropyLoss()

交叉熵损失非常适合多分类任务。这里有20种食品，所以模型需要在20个类别中选择正确答案。

优化器使用 AdamW：

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=0.0003,
    weight_decay=0.0001
)

优化器负责根据梯度更新模型参数。

lr=0.0003 是学习率，也就是每次参数调整的幅度。学习率过大，模型可能学得不稳定；学习率过小，训练速度又会比较慢。

weight_decay=0.0001 可以对过大的参数进行一定限制，降低模型过度记忆训练图片的风险。

十七、完整训练过程

代码一共训练10轮：

epochs = 10

每一轮先执行训练，再执行测试：

for epoch in range(1, epochs + 1):
    train(
        train_dataloader,
        model,
        loss_function,
        optimizer,
        epoch
    )

    test(
        test_dataloader,
        model,
        loss_function
    )

正常情况下，训练初期的损失会比较高，准确率也比较低。随着训练轮数增加，损失应该逐渐下降，测试准确率应该逐渐提高。

不过，训练轮数并不是越多越好。如果训练集表现越来越好，但测试集准确率开始下降，就可能出现了过拟合。也就是说，模型记住了训练图片，却没有真正掌握识别新图片的能力。

十八、总结

这个食品分类项目完成了从数据整理到模型保存的完整流程。

程序首先遍历食品图片目录，将数据划分为训练集和测试集；然后利用数据增强增加训练图片的变化；接着通过自定义 Dataset 和 DataLoader 分批读取数据；再使用三组卷积层提取图片特征，最后通过全连接层输出20种食品的预测分数。

在训练阶段，模型通过交叉熵计算错误，再利用 AdamW 优化器更新参数。每轮训练结束后，程序会使用测试集计算损失和准确率。如果当前准确率超过之前的最好结果，就自动保存新的最优模型。

整个流程可以总结为：

读取食品图片
生成训练集和测试集
对训练图片进行数据增强
分批送入卷积神经网络
提取图片特征
计算分类结果和损失
反向传播并更新参数
使用测试集计算准确率
保存效果最好的模型

这次练习让我认识到，一个图片分类模型并不是只靠几层神经网络就能完成。图片如何整理、数据如何处理、模型如何训练以及结果如何评价，每一步都很重要。

虽然目前做的是食品图片分类，但这套基本流程同样可以用在动物分类、垃圾分类和商品识别等任务中。只需要更换数据集、类别数量，并根据图片特点适当调整网络结构，就可以继续完成其他图片分类项目。