卷积神经网络搭建实战(二)——基于PyTorch框架和本地自定义图像数据集的用卷积神经网络实现的食物分类案例(全解三)
本文介绍了基于PyTorch框架的食物图像分类项目全流程。通过分析用户提供的5张数据集结构图,展示了标准的数据组织方式(train/test目录下按类别分子目录)。文章详细解析了代码实现,包括:1)自动生成数据路径索引文件;2)自定义Dataset类加载图像和标签;3)构建CNN网络模型;4)训练测试流程。实验预期准确率可达70-85%,并针对常见问题提供了优化建议。该项目适合PyTorch初学者
目录
引言
在计算机视觉领域,图像分类是核心技术之一,其应用场景覆盖医疗影像诊断、安防监控、智能推荐等多个领域。卷积神经网络(Convolutional Neural Network, CNN)凭借其局部感知和权值共享的特性,成为图像分类任务的“标配”模型。本文将以本地自定义食物图像数据集为对象,手把手演示如何使用PyTorch框架搭建卷积神经网络,完成从数据准备、模型构建到训练测试的全流程实战。
本文的所有代码和数据集结构均基于用户提供的实际项目素材(如图1-5所示),包含详细的注释和分步解析,适合PyTorch初学者及计算机视觉爱好者参考。
一、项目背景与数据集介绍
1.1 项目背景
食物分类是图像分类任务的垂直细分场景,其核心目标是通过模型识别图像中的食物种类(如薯条、八宝粥、骨肉相连等)。与通用图像分类(如ImageNet)不同,食物分类的数据集通常规模较小但类别更聚焦(本案例包含20类食物),因此需要更精细的数据组织和模型调优。
1.2 数据集结构(结合用户提供的5张图片)
用户提供的项目数据集结构清晰,完全符合PyTorch的Dataset
加载规范。通过图1-5的文件管理器界面,我们可以直观看到以下层级关系:
1.2.1 整体目录结构(图1、图2)
- 根路径:
2、卷积神经网络/data/食物分类/food_dataset
(图1、图2) - 子集文件夹:
train
(训练集)和test
(测试集),分别存储训练和验证用的图像数据(图2、图3)。 - 类别子文件夹:
train
和test
下各有多个以食物命名的子文件夹(如“八宝粥”“巴旦木”“薯条”等)(图3),每个子文件夹内存储对应类别的JPEG图像(图4、图5)。
这种“train
/test
目录下按类别分子目录”的结构是PyTorch官方推荐的自定义数据集标准结构,便于通过代码自动读取图像路径和标签。
图1 图2
图3 图4
图·5
1.2.2 具体类别与图像(图3、图4、图5)
- 类别示例:
train
文件夹下包含“八宝粥”“巴旦木”“白萝卜”“骨肉相连”“薯条”等20类食物(图3),每类食物对应一个独立子文件夹,test测试集文件结构完全相同,文件夹最底部图片数量较训练集少很多。 - 图像示例:以“八宝粥”类别为例(图4、图5),其文件夹内包含6张JPEG格式的图像(文件名如
img_八宝粥罐_22.jpeg
),每张图像均为食物特写,尺寸可能不统一(需通过代码统一缩放)。
这种结构下,标签可以通过“子文件夹名称”自动生成(如“薯条”文件夹的索引为n
,则该文件夹下所有图像的标签为n
),无需额外标注文件(如CSV或JSON),极大简化了数据预处理流程。
二、完整代码实现(附逐行注释)
2.1 环境准备与依赖导入
首先需要安装必要的库,包括PyTorch、Pillow(图像处理)、matplotlib(可视化)等。本文假设已配置好PyTorch环境(支持CUDA或MPS加速)。
# 导入基础库:用于文件操作、数值计算等
import os
import numpy as np
# 导入PyTorch核心库及数据加载工具
import torch
from torch.utils.data import Dataset, DataLoader # Dataset定义数据集,DataLoader批量加载数据
# 导入图像处理库:PIL用于打开、保存、显示图像
from PIL import Image
# 导入PyTorch的图像预处理工具(缩放、转Tensor等)
from torchvision import transforms
2.2 自动生成train.txt和test.txt文件
在PyTorch中,Dataset
类通常需要读取一个包含图像路径和标签的文本文件(如train.txt
)。本节通过遍历文件夹结构,自动生成这两个文件。
def generate_txt_files(root_dir, subset_dir, output_file):
"""
生成训练集或测试集的路径-标签列表文件(.txt)
参数:
root_dir (str): 数据集根目录(如'./食物分类/food_dataset')
subset_dir (str): 子集名称('train'或'test')
output_file (str): 输出文件路径(如'./train.txt')
"""
# 拼接子集的完整路径(如'./食物分类/food_dataset/train')
subset_path = os.path.join(root_dir, subset_dir)
# 打开输出文件(写入模式)
with open(output_file, 'w', encoding='utf-8') as f:
# 遍历子集路径下的所有目录和文件(os.walk递归遍历)
for root, dirs, files in os.walk(subset_path):
# 如果当前目录下有子目录(说明是类别文件夹的父级)
if dirs:
# 记录当前层级的所有类别名称(如['八宝粥', '巴旦木', ...])
class_names = dirs
else:
# 当前目录是类别文件夹(无嵌套子目录),获取类别名称(父目录的最后一级)
current_class = os.path.basename(root)
# 计算当前类别对应的标签(类别名称在class_names中的索引)
label = class_names.index(current_class)
# 遍历当前类别文件夹下的所有图像文件
for img_file in files:
# 拼接图像的完整路径(如'./食物分类/food_dataset/train/八宝粥/img_八宝粥罐_22.jpeg')
img_path = os.path.join(root, img_file)
# 将路径和标签写入txt文件(格式:"路径 标签")
f.write(f"{img_path} {label}
")
print(f"成功生成{subset_dir}集列表文件:{output_file}")
# 配置参数
dataset_root = r'.\食物分类\food_dataset' # 数据集根目录(根据实际路径修改)
train_subset = 'train' # 训练集子集名称
test_subset = 'test' # 测试集子集名称
# 生成训练集和测试集的txt文件
generate_txt_files(dataset_root, train_subset, 'train.txt')
generate_txt_files(dataset_root, test_subset, 'test.txt')
代码解析:
- 函数功能:
generate_txt_files
通过os.walk
递归遍历数据集目录,自动识别类别文件夹(如“八宝粥”),并将每个图像的路径与其对应的类别索引(标签)写入txt文件。 - 关键逻辑:
os.walk
返回三元组(root, dirs, files)
,其中root
是当前遍历的目录路径,dirs
是当前目录下的子目录名列表,files
是当前目录下的文件名列表。- 当
dirs
非空时,说明当前目录是类别文件夹的父级(如train
),此时记录所有子目录名(类别名称);当dirs
为空时,说明当前目录是类别文件夹(如train/八宝粥
),通过os.path.basename(root)
获取类别名称,并查找其在class_names
中的索引作为标签。 - 最终生成的
train.txt
和test.txt
格式示例如下:./食物分类/food_dataset/train/八宝粥/img_八宝粥罐_22.jpeg 0 ./食物分类/food_dataset/train/八宝粥/img_八宝粥罐_29.jpeg 0 ./食物分类/food_dataset/train/巴旦木/img_巴旦木_01.jpeg 1 ...
2.3 自定义数据集类(继承Dataset)
PyTorch的Dataset
类是一个抽象类,需要重写__len__
(数据集大小)和__getitem__
(按索引获取数据)方法。本节定义FoodDataset
类,用于加载图像和标签。
class FoodDataset(Dataset):
def __init__(self, txt_path, transform=None):
"""
初始化食物分类数据集
参数:
txt_path (str): 图像路径-标签列表文件的路径(如'train.txt')
transform (callable, optional): 图像预处理变换(如缩放、转Tensor)
"""
self.txt_path = txt_path
self.transform = transform
# 存储图像路径和标签的列表
self.img_paths = []
self.labels = []
# 读取txt文件并解析数据
with open(txt_path, 'r', encoding='utf-8') as f:
# 逐行读取,每行格式为"图像路径 标签"
for line in f.readlines():
# 去除行首尾空白,按空格分割(注意:路径可能包含空格,需谨慎处理)
parts = line.strip().split(' ')
if len(parts) == 2:
img_path = parts[0]
label = int(parts[1]) # 标签转为整数
self.img_paths.append(img_path)
self.labels.append(label)
def __len__(self):
"""返回数据集的大小(图像数量)"""
return len(self.img_paths)
def __getitem__(self, idx):
"""
按索引获取图像和标签
参数:
idx (int): 数据索引
返回:
tuple: (图像Tensor, 标签Tensor)
"""
# 获取图像路径和标签
img_path = self.img_paths[idx]
label = self.labels[idx]
# 用PIL打开图像(支持JPEG、PNG等格式)
image = Image.open(img_path).convert('RGB') # 确保图像为3通道(RGB)
# 应用预处理变换(如缩放、转Tensor)
if self.transform:
image = self.transform(image)
# 标签转为Tensor(PyTorch的损失函数需要Tensor类型)
label = torch.tensor(label, dtype=torch.long)
return image, label
代码解析:
- 初始化方法(
__init__
):读取txt
文件,将图像路径和标签分别存储到img_paths
和labels
列表中。注意处理路径中的空格问题(实际项目中建议避免路径包含空格)。 -
__len__
方法:返回数据集的大小,即图像的总数,通过len(self.img_paths)
获取。 -
__getitem__
方法:- 根据索引
idx
获取对应的图像路径和标签。 - 用
PIL.Image.open
打开图像,并通过.convert('RGB')
确保图像为3通道(避免灰度图导致的通道数不一致问题)。 - 应用预处理变换(如缩放、转Tensor),这部分在后续的
data_transforms
中定义。 - 将标签转为
long
类型的Tensor(PyTorch的交叉熵损失函数CrossEntropyLoss
要求标签为long
类型)。
- 根据索引
2.4 数据预处理与加载(DataLoader)
本节定义训练集和测试集的预处理变换,并通过DataLoader
批量加载数据。
# 定义训练集和测试集的预处理变换
data_transforms = {
'train': transforms.Compose([
transforms.Resize([256, 256]), # 调整图像大小为256x256(统一尺寸)
transforms.ToTensor(), # 转换为Tensor(自动归一化到[0,1])
]),
'valid':
transforms.Compose([
transforms.Resize([256, 256]),
transforms.ToTensor(),
]),
}#数组增强,
# 加载训练集和测试集
training_data = FoodDataset(
file_path = './train.txt', # 训练集的txt文件路径
transform = data_transforms['train'] # 应用训练集的预处理
)
test_data = FoodDataset(
file_path = './test.txt', # 测试集的txt文件路径
transform = data_transforms['valid'] # 应用测试集的预处理
)
# 创建数据加载器(DataLoader)
train_dataloader = DataLoader(training_data, batch_size=64,shuffle=True)#64张图片为一个包,
test_dataloader = DataLoader(test_data, batch_size=64,shuffle=True)
代码解析:
- 预处理变换(
data_transforms
):transforms.Resize([256, 256])
:将图像尺寸统一为256x256,确保所有图像能组成批量Tensor(若图像尺寸不一致,无法进行批量运算)。transforms.ToTensor()
:将PIL图像转换为PyTorch的Tensor,并自动将像素值从[0, 255]
归一化到[0, 1]
。
- DataLoader:
batch_size=64
:每批加载64张图像,平衡内存占用和训练效率。shuffle=True
(训练集):打乱数据顺序,避免模型因数据顺序固定而学习到无关模式;测试集shuffle=False
,便于按顺序评估结果。
2.5 模型构建(CNN网络定义)
本节定义一个简单的卷积神经网络(CNN),包含3个卷积层和2个全连接层,适用于食物分类任务。
import torch.nn as nn
class CNN(nn.Module):
def __init__(self): # 输入大小 (1,3,256,256)(1,3,288,288)
super(CNN, self).__init__()#
self.conv1 = nn.Sequential( #将多个层组合成一起。
nn.Conv2d( #2d一般用于图像,3d用于视频数据(多一个时间维度),1d一般用于结构化的序列数据
in_channels=3, # 图像通道个数,1表示灰度图(确定了卷积核 组中的个数),3表示RGB彩色图像
out_channels=16,# 要得到几多少个特征图,卷积核的个数,
kernel_size=5, # 卷积核大小,5 * 5
stride=1, # 步长
padding=2, # 一般希望卷积核处理后的结果大小与处理前的数据大小相同,效果会比较好。那padding改如何设计呢?建议stride为1,kernel_size = 2*padding+1
), # 输出的特征图为 (1,16,256,256)(1,16,288,288)
nn.ReLU(), # relu层(1,16,256,256)(1,16,308,308)
nn.MaxPool2d(kernel_size=2), # 进行池化操作(2x2 区域), 输出结果为:(1,16,128,128)(1,16,144,144)
)
self.conv2 = nn.Sequential( #输入
nn.Conv2d(16, 32, 5, 1, 2), # 输出 (1,32,128,128) (1,32,144,144)
nn.ReLU(), # relu层(1,32,128,128) (1,32,144,144)
nn.Conv2d(32, 32, 5, 1, 2), # 输出(1,32,128,128)(1,32,144,144)
nn.ReLU(),#输出(1,32,128,128)(1,32,144,144)
nn.MaxPool2d(2), # 输出 (1,32,64,64) (1,32,72,72)
)
self.conv3 = nn.Sequential( #输入 (32,64,64)
nn.Conv2d(32, 128, 5, 1, 2), #(1,128,64,64) (1,128,72,72)
nn.ReLU(), # 输出(1,128,64,64) (1,128,72,72)
)
self.out = nn.Linear(128 * 64 * 64,20) # 全连接层得到的结果 (128 * 63 * 44,20)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)# 输出 (64,128, 64, 64)
x = x.view(x.size(0), -1) # flatten操作,结果为:(batch_size, 64 * 32 * 32)
output = self.out(x)
return output
model = CNN().to(device)
print(model)
代码解析:
- 卷积层设计:
nn.Conv2d
:二维卷积层,用于提取图像的空间特征。in_channels
为输入通道数(3表示RGB图像),out_channels
为输出通道数(即卷积核数量,生成对应数量的特征图),kernel_size
为卷积核尺寸(5x5),stride
为步长(1表示每次滑动1个像素),padding
为填充(2表示在图像边缘填充2个像素,确保输出尺寸与输入一致)。nn.ReLU()
:修正线性单元激活函数,引入非线性特性,使模型能拟合更复杂的特征。nn.MaxPool2d
:最大池化层,通过取局部区域的最大值降低特征图尺寸(宽、高减半),减少计算量的同时保留主要特征。
- 全连接层:
nn.Linear(128 * 64 * 64,20)
:将卷积后的特征图展平为一维向量(长度为128 * 64 * 64
),然后通过全连接层映射到20维的输出(对应20类食物的分类得分)。
2.6 模型训练与测试
本节定义训练函数和测试函数,完成模型的优化和性能评估。
def train(dataloader, model, loss_fn, optimizer):
model.train() #放开w参数的修改权限
#pytorch提供2种方式来切换训练和测试的模式,分别是:model.train() 和 model.eval()。
# 一般用法是:在训练开始之前写上model.trian(),在测试时写上 model.eval() 。
batch_size_num = 1
for X, y in dataloader: #其中batch为每一个数据的编号
X, y = X.to(device), y.to(device) #把训练数据集和标签传入cpu或GPU
pred = model.forward(X) #自动初始化 w权值
loss = loss_fn(pred, y) #通过交叉熵损失函数计算损失值loss
# Backpropagation 进来一个batch的数据,计算一次梯度,更新一次网络
optimizer.zero_grad() #梯度值清零
loss.backward() #反向传播计算得到每个参数的梯度值
optimizer.step() #根据梯度更新网络参数
loss = loss.item() #获取损失值
if batch_size_num %1 == 0:
print(f"loss: {loss:>7f} [number:{batch_size_num}]")
batch_size_num += 1
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset)
num_batches = len(dataloader)
model.eval() #测试模式
test_loss, correct = 0, 0
with torch.no_grad(): #一个上下文管理器,关闭梯度计算。当你确认不会调用Tensor.backward()的时候。这可以减少计算所用内存消耗。
for X, y in dataloader:
X, y = X.to(device), y.to(device)
pred = model.forward(X)
test_loss += loss_fn(pred, y).item() #
correct += (pred.argmax(1) == y).type(torch.float).sum().item()
a = (pred.argmax(1) == y) #dim=1表示每一行中的最大值对应的索引号,dim=0表示每一列中的最大值对应的索引号
b = (pred.argmax(1) == y).type(torch.float)
test_loss /= num_batches
correct /= size
print(f"Test result: \n Accuracy: {(100*correct)}%, Avg loss: {test_loss}")
loss_fn = nn.CrossEntropyLoss() #创建交叉熵损失函数对象,因为食物的类别是20,
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)#创建一个优化器,SGD为随机梯度下降算法??
# params:要训练的参数,一般我们传入的都是model.parameters()。
# lr:learning_rate学习率,也就是步长。
# train(train_dataloader, model, loss_fn, optimizer)
# test(test_dataloader, model, loss_fn)
epochs = 10
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
# test(test_dataloader, model, loss_fn)
print("Done!")
test(test_dataloader, model, loss_fn)
代码解析:
- 训练函数(
train
):model.train()
:开启训练模式,启用Dropout层(随机失活神经元)和BatchNorm层(计算当前批次的均值和方差)。- 遍历
dataloader
获取每个批次的图像(X
)和标签(y
),将数据和标签移动到目标设备(CPU/GPU)。 - 前向传播计算预测值
pred
,通过损失函数loss_fn
计算预测值与真实标签的损失loss
。 - 反向传播
loss.backward()
计算梯度,优化器optimizer.step()
更新模型参数。 - 统计每个批次的损失值,按指定频率打印训练进度。
- 测试函数(
test
):model.eval()
:开启测试模式,禁用Dropout和BatchNorm的随机操作(使用训练阶段统计的均值和方差)。with torch.no_grad()
:关闭自动梯度计算,减少内存消耗。- 遍历测试集,计算整体损失和准确率,评估模型的泛化能力。
2.7 模型初始化与训练执行
最后,初始化模型、损失函数和优化器,启动训练流程。
if __name__ == '__main__':
# 配置参数
epochs = 10 # 训练轮数
# 自动选择计算设备(优先GPU,其次Apple 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 = CNN().to(device)
loss_fn = nn.CrossEntropyLoss() # 交叉熵损失函数(适用于多分类)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Adam优化器
# 启动训练
print("开始训练...")
for t in range(epochs):
print(f"Epoch {t+1}\n-------------------------------")
train(train_dataloader, model, loss_fn, optimizer)
# 测试训练后的模型
print("
训练结束,开始测试...")
test(test_dataloader, model, loss_fn)
代码解析:
- 设备选择:通过条件判断自动检测可用的计算设备(GPU优先,其次Apple MPS,最后CPU),提升训练速度。
- 损失函数:
nn.CrossEntropyLoss
结合了Softmax激活函数和交叉熵损失,适用于多分类任务(输出为类别得分,无需手动计算概率)。 - 优化器:
torch.optim.Adam
是一种自适应学习率的优化器,通常比传统的SGD收敛更快。
三、代码分段深度解析
3.1 数据准备:为什么选择“按类别分文件夹”?
用户提供的图片(图1-5)展示了标准的“train
/test
目录下按类别分子目录”结构,这是PyTorch的Dataset
加载机制的最佳实践,主要原因如下:
- 自动化标签生成:通过遍历子目录名(如“八宝粥”“薯条”),代码自动生成标签(子目录索引),无需手动标注。
- 兼容数据增强:在
Dataset
的__getitem__
方法中应用数据增强(如随机翻转),可以灵活地对每个批次的图像进行变换,避免重复存储增强后的图像。 - 符合官方最佳实践:PyTorch官方教程(如
ImageFolder
)推荐此结构,社区支持广泛,便于后续扩展(如使用预训练模型的ImageFolder
加载方式)。
3.2 自定义Dataset类的核心方法
FoodDataset
类继承自nn.Module
的Dataset
,其核心是__len__
和__getitem__
:
__len__
:告诉PyTorch数据集有多大,DataLoader
通过此方法确定需要生成多少个批次。__getitem__
:定义如何获取单个样本。这里不仅返回图像和标签,还应用了预处理变换(如缩放、转Tensor),确保数据在输入模型前格式正确。
注意:如果图像尺寸不一致,Resize
变换是必须的(否则无法组成批量Tensor);如果数据量极大,可使用num_workers>0
加速数据加载(但需注意Windows系统的多线程问题)。
3.3 CNN模型的设计逻辑
本文的CNN
是一个典型的“卷积-池化-全连接”结构,设计逻辑如下:
- 浅层卷积(conv1):提取边缘、纹理等低级特征(如食物的轮廓、颜色块)。
- 中层卷积(conv2):组合低级特征形成中级特征(如食物的局部结构,如薯条的条状、八宝粥的颗粒感)。
- 深层卷积(conv3):提取高级语义特征(如食物的整体形状、细节,如骨肉相连的骨头和肉的纹理)。
- 全连接层:将高级特征映射到类别空间,输出每个类别的得分。
改进方向:实际项目中可使用更深的网络(如ResNet、VGG),或引入残差连接、注意力机制(如SE Block)提升性能。
3.4 训练过程中的关键细节
- 数据归一化:
ToTensor
变换将像素值从[0, 255]
归一化到[0, 1]
,Normalize
变换(若添加)进一步将数据标准化到均值为0、标准差为1的分布,有助于优化器更快收敛。 - 学习率调整:本文使用固定学习率(0.001),实际中可使用
torch.optim.lr_scheduler
动态调整(如余弦退火、阶梯衰减)。 - 过拟合处理:除了数据增强,还可添加Dropout层、L2正则化(
weight_decay
参数)等方法。
四、实验结果与优化建议
4.1 预期结果
假设数据集包含20类食物,每类约100张图像(训练集80%,测试集20%),使用上述代码训练10轮后,预期测试准确率可达70%-85%(具体取决于数据质量和类别复杂度)。
4.2 常见问题与优化建议
-
问题1:训练损失不下降
- 可能原因:学习率过大(模型无法收敛)或过小(收敛过慢)、数据预处理错误(如标签错误)、模型容量不足(网络太浅)。
- 解决方法:调整学习率(如从0.001降至0.0001)、检查
txt
文件中的标签是否正确、增加卷积层或全连接层的神经元数量。
-
问题2:测试准确率远低于训练准确率
- 可能原因:过拟合(模型过度记忆训练数据)。
- 解决方法:增加数据增强(如
RandomHorizontalFlip
、RandomRotation
)、添加Dropout层(如在conv1
后加nn.Dropout2d(0.5)
)、使用早停法(Early Stopping)。
-
问题3:GPU内存不足
- 可能原因:批量大小(
batch_size
)过大、图像尺寸过大。 - 解决方法:减小
batch_size
(如从64降至32)、降低图像尺寸(如从256x256降至128x128)、使用混合精度训练(torch.cuda.amp
)。
- 可能原因:批量大小(
五、总结
本文通过完整的代码示例和详细解析,演示了如何基于PyTorch框架和本地自定义数据集实现食物分类任务。核心步骤包括:
- 数据准备:按类别分文件夹组织数据,自动生成
train.txt
和test.txt
。 - 自定义数据集:继承
Dataset
类,实现__len__
和__getitem__
方法。 - 模型构建:设计卷积神经网络,提取图像特征并映射到类别空间。
- 训练与测试:定义损失函数和优化器,通过
DataLoader
批量加载数据,迭代训练并评估模型性能。
通过本案例,读者可以掌握PyTorch处理自定义数据集的核心流程,并为后续的图像分类、目标检测等任务打下坚实基础。建议在实际项目中尝试不同的网络结构(如ResNet)、数据增强方法(如RandomCrop
)和超参数调优(如学习率),以进一步提升模型性能。
更多推荐
所有评论(0)