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

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

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

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

python 复制代码
# 文件操作相关库
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 编写路径整理函数

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

python 复制代码
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 定义数据变换

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

python 复制代码
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提供可迭代的数据接口。

python 复制代码
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生成的数据按批次划分,并支持打乱数据顺序,为模型训练提供高效的批量数据迭代器。

python 复制代码
# 实例化训练集和测试集
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 个全连接层,用于提取图像特征并完成分类任务。

python 复制代码
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 定义训练函数

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

python 复制代码
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 定义测试函数

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

python 复制代码
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 执行训练与评估

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

python 复制代码
# 定义损失函数:交叉熵损失(适用于多分类任务)
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)

六、单张图像预测

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

python 复制代码
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. 若需进一步提升性能,可增加数据增强操作(如随机裁剪、翻转、旋转等),并使用预训练模型进行迁移学习。
相关推荐
weixin_456904277 小时前
深度学习模型边缘部署与B/S架构
人工智能·深度学习·架构
MichaelIp7 小时前
利用ms-swift微调和百炼平台微调大模型
人工智能·gpt·自然语言处理·prompt·aigc·swift·agi
SHIPKING3937 小时前
【机器学习&深度学习】向量模型与重排序模型:RAG 的双引擎解析
人工智能·深度学习·机器学习·向量模型·重排序模型
飞哥数智坊7 小时前
元宝AI:微信中帮我看阅兵的“朋友”
人工智能
lypzcgf8 小时前
Coze源码分析-工作空间-资源查询-后端源码
人工智能·后端·系统架构·开源·go
海天一色y8 小时前
动手学深度学习
人工智能·深度学习
恒点虚拟仿真8 小时前
AI+虚拟仿真实践教学,照亮智慧能源教育之路
人工智能·ai·能源·虚拟仿真实验·人工智能+虚拟仿真·ai教学
fzy00858 小时前
教育项目管理工具新趋势:可视化与自动化如何提升效率?
大数据·人工智能·自动化
AI Echoes8 小时前
LangGraph 重要注意事项和常见问题
人工智能·python·langchain·agent