用 PyTorch 实现食品图像分类：从数据预处理到模型训练与预测

在计算机视觉领域，图像分类是基础且重要的任务之一。本文将详细介绍如何使用 PyTorch 框架，完成一套完整的食品图像分类流程，包括数据路径整理、自定义数据集构建、卷积神经网络模型搭建、模型训练与评估，以及最终的单张图片预测功能。

一、前期准备：导入所需库

首先，我们需要导入实验过程中用到的各类库，涵盖文件操作、数据处理、模型构建与训练等多个方面。

# 文件操作相关库import os# 深度学习框架核心库import torchfrom torch.utils.data import Dataset, DataLoaderfrom torch import nnfrom torchvision import transforms# 数据处理与图像读取库import numpy as npfrom PIL import Image

二、数据预处理：整理图像路径与标签

在进行模型训练前，我们需要将分散在不同文件夹中的图像数据，整理成 “图像路径 - 类别标签” 的对应关系，并保存到文本文件中，方便后续数据集读取。

2.1 编写路径整理函数

该函数会遍历指定文件夹下的所有图像文件，获取每个图像的完整路径，并根据其所在文件夹名称分配类别标签，最后将结果写入文本文件。

dirs = [] # 用于存储类别名称，索引对应类别标签def train_test_file(root, dir): # 以写入模式创建文本文件，用于保存图像路径和标签 file_txt = open(dir + \'.txt\', \'w\') # 拼接完整的目标文件夹路径 path = os.path.join(root, dir) # 遍历目标文件夹下的所有子文件夹和文件 for roots, directories, files in os.walk(path): # 若当前目录存在子文件夹，将子文件夹名称存入dirs（即类别名称） if len(directories) != 0: for i in directories: dirs.append(i) # 若当前目录为图像文件所在目录（无更多子文件夹） else: # 分割路径，获取当前文件夹名称（即类别名称） now_dir = roots.split(\'\\\\\') # 遍历当前目录下的所有图像文件 for file in files: # 拼接图像文件的完整路径 path_1 = os.path.join(roots, file) # 将图像路径和对应类别标签写入文本文件 file_txt.write(path_1 + \' \' + str(dirs.index(now_dir[-1])) + \'\\n\') # 关闭文本文件 file_txt.close()# 定义数据集根路径和训练/测试文件夹名称root = r\"D:\\pythonProject11\\深度学习\\目录1\\food_dataset\"train_dir = \'train\'test_dir = \'test\'# 分别处理训练集和测试集，生成对应的路径-标签文本文件train_test_file(root, train_dir)train_test_file(root, test_dir)

三、构建自定义数据集

PyTorch 提供了Dataset抽象类，我们通过继承该类，实现自定义的食品图像数据集，支持图像读取、数据增强与标签处理。

3.1 定义数据变换

为了提升模型泛化能力并统一输入格式，我们对训练集和验证集（测试集）分别定义数据变换操作，主要包括图像尺寸调整和张量转换。

data_transform = { \'train\': transforms.Compose([ transforms.Resize([256, 256]), # 将训练集图像调整为256×256尺寸 transforms.ToTensor(), # 将PIL图像转换为PyTorch张量（维度：C×H×W） ]), \'valid\': transforms.Compose([ transforms.Resize([256, 256]), # 测试集图像同样调整为256×256尺寸 transforms.ToTensor(), ]),}

3.2 实现自定义 Dataset 类

该类会从之前生成的文本文件中读取图像路径和标签，在__getitem__方法中完成图像读取、变换与标签转换，为后续DataLoader提供可迭代的数据接口。

class food_dataset(Dataset): # 继承Dataset抽象类 def __init__(self, file_path, transform=None): self.file_path = file_path # 存储路径-标签文本文件的路径 self.imgs = [] # 存储所有图像的路径 self.label = [] # 存储所有图像对应的标签 self.transform = transform # 存储数据变换方法 # 读取文本文件，解析图像路径和标签 with open(self.file_path) as f: # 按行读取文件内容，分割路径和标签 samples = [x.strip().split(\' \') for x in f.readlines()] for img_path, label in samples: self.imgs.append(img_path) self.label.append(label) # 返回数据集的总样本数 def __len__(self): return len(self.imgs) # 根据索引获取单个样本（图像+标签） def __getitem__(self, idx): # 读取图像（默认以PIL格式打开） image = Image.open(self.imgs[idx]) # 若存在数据变换，对图像进行处理 if self.transform: image = self.transform(image) # 将标签转换为PyTorch张量（int64类型，适配交叉熵损失） label = self.label[idx] label = torch.from_numpy(np.array(label, dtype=np.int64)) return image, label

3.3 构建 DataLoader

DataLoader会将Dataset生成的数据按批次划分，并支持打乱数据顺序，为模型训练提供高效的批量数据迭代器。

# 实例化训练集和测试集training_data = food_dataset(file_path=\'train.txt\', transform=data_transform[\'train\'])test_data = food_dataset(file_path=\'test.txt\', transform=data_transform[\'valid\'])# 构建训练集和测试集的DataLoadertrain_dataloader = DataLoader(training_data, batch_size=2, shuffle=True) # 批次大小2，训练集打乱test_dataloader = DataLoader(test_data, batch_size=2, shuffle=True) # 批次大小2，测试集打乱

四、搭建卷积神经网络模型

我们设计一个简单的卷积神经网络（CNN），包含 3 个卷积层（含激活函数和池化层）和 1 个全连接层，用于提取图像特征并完成分类任务。

class Sequentialnetwork(nn.Module): def __init__(self): super().__init__() # 调用父类nn.Module的构造函数 # 第一个卷积块：卷积层+ReLU激活函数+最大池化层 self.conv1 = nn.Sequential( nn.Conv2d( in_channels=3, # 输入通道数：RGB图像为3 out_channels=16, # 输出通道数：16个卷积核 kernel_size=5, # 卷积核大小：5×5 stride=1, # 步长：1 padding=2, # 填充：2（保证输入输出尺寸一致） ), nn.ReLU(), # ReLU激活函数，引入非线性 nn.MaxPool2d(kernel_size=2) # 最大池化层：2×2池化核，尺寸减半 ) # 第二个卷积块：卷积层+ReLU激活函数（无池化层） self.conv2 = nn.Sequential( nn.Conv2d(16, 64, 5, 1, 2), # 输入16通道，输出64通道，其他参数同上 nn.ReLU(), ) # 第三个卷积块：卷积层+ReLU激活函数+最大池化层 self.conv3 = nn.Sequential( nn.Conv2d(64, 128, 5, 1, 2), # 输入64通道，输出128通道 nn.ReLU(), nn.MaxPool2d(kernel_size=2) # 池化后尺寸再次减半 ) # 全连接层：将卷积提取的特征映射为类别概率（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) # 展平特征图：(批次大小, 通道数×高×宽) x = self.out(x) # 经过全连接层，输出类别得分 return x# 检测并选择训练设备（优先GPU，其次MPS，最后CPU）device = \'cuda\' if torch.cuda.is_available() else \'mps\' if torch.backends.mps.is_available() else \'cpu\'print(f\'Using {device} device\')# 实例化模型并将其移动到指定设备model = Sequentialnetwork().to(device)print(model) # 打印模型结构，查看各层参数

五、模型训练与评估

5.1 定义训练函数

训练函数负责模型的迭代训练过程，包括前向传播计算损失、反向传播更新梯度，并定期打印训练损失。

def train(dataloader, model, loss_fn, optimizer): model.train() # 切换模型为训练模式（启用Dropout、BatchNorm等训练特有的层） batch_size_num = 1 # 记录当前训练的批次号 # 遍历DataLoader中的每一个批次 for X, y in dataloader: # 将数据移动到指定设备（与模型设备一致） X, y = X.to(device), y.to(device) # 前向传播：计算模型预测值 pre = model.forward(X) # 计算损失（交叉熵损失，适用于多分类任务） loss = loss_fn(pre, y) # 反向传播与参数更新 optimizer.zero_grad() # 清空上一轮的梯度（避免梯度累积） loss.backward() # 反向传播，计算各参数的梯度 optimizer.step() # 根据梯度更新模型参数 # 获取当前批次的损失值（转换为Python数值） loss_value = loss.item() # 每100个批次打印一次损失值，监控训练进度 if batch_size_num % 100 == 0: print(f\"loss:{loss_value:>7f} [number:{batch_size_num}]\") batch_size_num += 1

5.2 定义测试函数

测试函数用于评估模型在测试集上的性能，包括计算平均损失和分类准确率，判断模型的泛化能力。

def test(dataloader, model, loss_fn): model.eval() # 切换模型为评估模式（禁用Dropout、固定BatchNorm参数） size = len(dataloader.dataset) # 测试集总样本数 num_batches = len(dataloader) # 测试集总批次数 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() # 统计正确预测数：预测类别（argmax(1)取得分最高的类别）与真实标签一致的数量 correct += (pred.argmax(1) == y).type(torch.float).sum().item() # 计算测试集的平均损失和准确率 test_loss /= num_batches correct /= size # 打印测试结果 print(f\"Test Error: \\n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")

5.3 执行训练与评估

设置损失函数、优化器和训练轮次，然后执行训练流程，并在训练结束后评估模型在测试集上的性能。

# 定义损失函数：交叉熵损失（适用于多分类任务）loss_fn = nn.CrossEntropyLoss()# 定义优化器：Adam优化器（学习率0.01）optimizer = torch.optim.Adam(model.parameters(), lr=0.01)# 先进行一轮预热训练train(train_dataloader, model, loss_fn, optimizer)# 定义总训练轮次epochs = 10# 迭代训练for t in range(epochs): print(f\'Epoch {t+1}\\n-------------------------------\') train(train_dataloader, model, loss_fn, optimizer)print(\'Training Done!\')# 训练结束后，在测试集上评估模型性能test(test_dataloader, model, loss_fn)

六、单张图像预测

训练完成后，我们可以使用训练好的模型对单张食品图像进行分类预测，输出其类别。

def predict(img_path, model): model.eval() # 切换模型为评估模式 # 读取图像并转换为RGB格式（避免灰度图或其他格式导致通道数不匹配） image = Image.open(img_path).convert(\'RGB\') # 对图像进行预处理（与测试集一致），并增加批次维度（模型输入需为4维：批次×通道×高×宽） X = data_transform[\'valid\'](image).unsqueeze(0).to(device) # 禁用梯度计算 with torch.no_grad(): pred = model(X) # 前向传播获取预测结果 # 找到得分最高的类别索引，并根据dirs列表获取类别名称 pred_class_idx = pred.argmax(1).item() print(f\"预测结果: {dirs[pred_class_idx]}\")# 接收用户输入的图像路径，并进行预测img_path = input(\"请输入要预测的图像路径：\")predict(img_path, model)

七、总结与说明

本文实现了一套完整的食品图像分类流程，从数据预处理到模型训练、评估与预测，涵盖了 PyTorch 计算机视觉项目的核心环节。需要注意的是：

1. 数据路径需根据实际情况修改（root变量），确保文本文件能正确生成。
2. 模型结构可根据数据集复杂度调整（如增加卷积层、调整通道数、添加 Dropout 层等），以提升性能。
3. 学习率、批次大小、训练轮次等超参数需根据实验效果优化，避免过拟合或训练缓慢。
4. 若需进一步提升性能，可增加数据增强操作（如随机裁剪、翻转、旋转等），并使用预训练模型进行迁移学习。