在计算机视觉领域，图像分类是基础且重要的任务之一。本文将详细介绍如何使用 PyTorch 框架，完成一套完整的食品图像分类流程，包括数据路径整理、自定义数据集构建、卷积神经网络模型搭建、模型训练与评估，以及最终的单张图片预测功能。
一、前期准备：导入所需库
首先，我们需要导入实验过程中用到的各类库，涵盖文件操作、数据处理、模型构建与训练等多个方面。
# 文件操作相关库
import os
# 深度学习框架核心库
import torch
from torch.utils.data import Dataset, DataLoader
from torch import nn
from torchvision import transforms
# 数据处理与图像读取库
import numpy as np
from 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'])
# 构建训练集和测试集的DataLoader
train_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 计算机视觉项目的核心环节。需要注意的是：
- 数据路径需根据实际情况修改（
root变量），确保文本文件能正确生成。 - 模型结构可根据数据集复杂度调整（如增加卷积层、调整通道数、添加 Dropout 层等），以提升性能。
- 学习率、批次大小、训练轮次等超参数需根据实验效果优化，避免过拟合或训练缓慢。
- 若需进一步提升性能，可增加数据增强操作（如随机裁剪、翻转、旋转等），并使用预训练模型进行迁移学习。