1. 项目概述
在这个项目中,我们将使用PyTorch框架构建一个卷积神经网络(CNN)来实现食物图像分类任务。我们的数据集包含20种不同的食物类别,包括八宝粥、巴旦木、白萝卜、板栗等常见食物。本文将详细介绍从数据准备、模型构建到训练和评估的完整流程。
2. 环境准备
首先,我们需要导入必要的Python库:
python
import torch
import torchvision.models as models
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np
import os
3. 数据准备与预处理
3.1 数据转换定义
我们定义了两组数据转换,分别用于训练集和验证集:
python
data_transforms = {
'train': transforms.Compose([
transforms.Resize([256,256]), # 统一图像大小为256x256
transforms.ToTensor(), # 转换为PyTorch张量
]),
'valid': transforms.Compose([
transforms.Resize([256,256]),
transforms.ToTensor(),
]),
}
3.2 数据文件准备
我们编写了一个函数来生成包含图像路径和标签的文本文件:
python
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):
if len(directories) != 0:
dirs = directories
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'.\食物分类\food_dataset'
train_dir = 'train'
test_dir = 'test'
train_test_file(root, train_dir)
train_test_file(root, test_dir)
- 生成的 txt 文件如下所示,包含图片的路径、类别和对应的标签

3.3 自定义数据集类
我们创建了一个继承自torch.utils.data.Dataset
的自定义数据集类:
python
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().split(' ') for x in f.readlines()]
for img_path, label in samples:
self.imgs.append(img_path)
self.labels.append(label)
def __len__(self):
return len(self.imgs)
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
3.4 创建数据加载器
python
training_data = food_dataset(file_path='train.txt', transform=data_transforms['train'])
test_data = food_dataset(file_path='test.txt', transform=data_transforms['valid'])
train_dataloader = DataLoader(training_data, batch_size=16, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=16, shuffle=True)
4. 模型构建
我们定义了一个包含三个卷积层的CNN模型:
python
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
# 第一个卷积块
self.conv1 = nn.Sequential(
nn.Conv2d(3, 16, 5, 1, 2), # (16,256,256)
nn.ReLU(),
nn.MaxPool2d(2), # (16,128,128)
)
# 第二个卷积块
self.conv2 = nn.Sequential(
nn.Conv2d(16, 32, 5, 1, 2), # (32,128,128)
nn.ReLU(),
nn.MaxPool2d(2), # (32,64,64)
)
# 第三个卷积块
self.conv3 = nn.Sequential(
nn.Conv2d(32, 64, 5, 1, 2), # (64,64,64)
nn.ReLU(),
nn.MaxPool2d(2), # (64,32,32)
)
# 全连接层
self.out = nn.Linear(64*32*32, 20)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1) # 展平操作
output = self.out(x)
return output