从零搭建卷积神经网络(CNN):食品图像分类实战

卷积神经网络(CNN)是深度学习领域处理图像任务的核心模型,凭借对图像空间特征的高效提取能力,在图像分类、目标检测等任务中表现卓越。本文将以食品图像分类为例,从零讲解如何用 PyTorch 搭建完整的 CNN 训练流程,涵盖数据集构建、网络设计、模型训练与评估全环节。

一、项目背景与技术栈

1. 项目目标

基于自定义食品数据集(20 类食品),搭建 CNN 模型实现图像分类,输入为 256×256 的 RGB 图像,输出为食品类别预测结果。

2. 核心技术栈

  • 框架:PyTorch(简洁的动态图机制,适合新手入门)
  • 数据处理:PIL(图像读取)、torchvision.transforms(图像预处理)
  • 模型核心:卷积层、池化层、全连接层
  • 优化策略:Adam 优化器 + 交叉熵损失函数

二、完整实现流程

步骤 1:环境准备与模块导入

首先导入所需的核心模块,涵盖数据加载、模型构建、图像预处理等功能:

python 复制代码
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
from torchvision import transforms

步骤 2:自定义数据集构建

PyTorch 通过Dataset抽象类实现自定义数据集,核心需实现__init__、len、__getitem__三个方法:

python 复制代码
# 1. 图像预处理:统一尺寸+转为张量
data_transforms = {
    'trainda': transforms.Compose([
        transforms.Resize([256, 256]),  # 统一图像尺寸为256×256
        transforms.ToTensor(),         # 转为张量(通道×高度×宽度)
    ]),
    'valid': transforms.Compose([
        transforms.Resize([256, 256]),
        transforms.ToTensor(),
    ]),
}

# 2. 自定义Dataset类
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):
        # 读取图像(默认RGB格式)
        image = Image.open(self.imgs[idx]).convert('RGB')
        # 应用预处理
        if self.transform:
            image = self.transform(image)
        # 标签转为64位整数张量(适配CrossEntropyLoss)
        label = torch.from_numpy(np.array(self.labels[idx], dtype=np.int64))
        return image, label

# 3. 实例化数据集并创建数据加载器
training_data = food_dataset(file_path='train.txt', transform=data_transforms['trainda'])
test_data = food_dataset(file_path='test.txt', transform=data_transforms['valid'])

# DataLoader实现批量加载+打乱数据
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)
关键说明:
  • 标签文件格式:train.txt/test.txt每行需满足「图像路径 数字标签」(如/apple_1.jpg 0);
  • ToTensor()会将图像像素值从 [0,255] 归一化到 [0,1],并转换为(C,H,W)格式(PyTorch 标准);
  • convert('RGB')确保灰度图 / 异常通道图像转为 3 通道,避免维度不匹配。

步骤 3:卷积神经网络设计

CNN 的核心是通过卷积层提取空间特征,池化层降维,全连接层完成分类。本文设计的网络结构如下:

python 复制代码
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 卷积层1:3→16通道,5×5卷积核,步长1,填充2(保持尺寸不变)
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),  # 激活函数:引入非线性
            nn.MaxPool2d(kernel_size=2),  # 最大池化:尺寸减半(256→128)
        )
        # 卷积层2:16→32通道,连续2次卷积+池化(128→64)
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # 卷积层3:32→128通道(无池化,保留64×64尺寸)
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=128, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
        )
        # 全连接层:展平特征→20类分类(20为食品类别数)
        self.out = nn.Linear(128*64*64, out_features=20)

    # 前向传播(定义数据流动路径)
    def forward(self, x):
        x = self.conv1(x)   # 输出:(batch, 16, 128, 128)
        x = self.conv2(x)   # 输出:(batch, 32, 64, 64)
        x = self.conv3(x)   # 输出:(batch, 128, 64, 64)
        x = x.view(x.size(0), -1)  # 展平:(batch, 128*64*64)
        output = self.out(x)       # 输出:(batch, 20)
        return output

# 设备选择:优先GPU(CUDA),否则CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)  # 模型部署到指定设备
print(model)  # 打印网络结构
核心知识点:
  1. 卷积层参数设计:填充padding=(kernel_size-1)//2时,卷积后图像尺寸不变(如 5×5 卷积核对应 padding=2);
  2. 池化层作用:最大池化(MaxPool2d)保留关键特征,同时将尺寸减半,降低计算量;
  3. 展平操作:x.view(x.size(0), -1)将 4 维特征张量(batch, C, H, W)转为 2 维(batch, C×H×W),适配全连接层输入;
  4. 设备部署model.to(device)将模型参数迁移到 GPU/CPU,需保证后续数据与模型在同一设备。

步骤 4:训练与评估函数实现

1. 训练函数(trainda)

负责单轮训练的前向传播、反向传播与参数更新:

python 复制代码
def trainda(dataloader, model, loss_fn, optimizer):
    model.train()  # 训练模式:启用Dropout/BatchNorm等层的训练行为
    batch_size_num = 1
    for X, Y in dataloader:
        # 数据部署到指定设备
        X, Y = X.to(device), Y.to(device)
        
        # 1. 前向传播:预测结果
        pred = model(X)  # 等价于model.forward(X)
        # 2. 计算损失(交叉熵损失)
        loss = loss_fn(pred, Y)
        
        # 3. 反向传播:梯度清零→计算梯度→更新参数
        optimizer.zero_grad()  # 清空历史梯度(避免累积)
        loss.backward()        # 反向传播计算梯度
        optimizer.step()       # 优化器更新参数
        
        # 打印批次损失
        loss_val = loss.item()  # 张量转标量
        if batch_size_num % 10 == 0:  # 每10批次打印一次(优化原逻辑冗余)
            print(f'loss: {loss_val:>7f} [batch:{batch_size_num}]')
        batch_size_num += 1
2. 评估函数(testda)

在测试集上评估模型性能(无梯度计算,节省内存):

python 复制代码
def testda(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()  # 评估模式:关闭Dropout/BatchNorm的训练行为
    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()
            # 预测类别:取概率最大的索引(dim=1表示按行取)
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    
    # 计算平均损失与准确率
    test_loss /= num_batches
    correct /= size
    print(f"Test Result: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f}")

步骤 5:模型训练主循环

设置损失函数、优化器,执行多轮训练与评估:

python 复制代码
# 1. 损失函数:交叉熵损失(适配分类任务)
loss_fn = nn.CrossEntropyLoss()
# 2. 优化器:Adam(自适应学习率,收敛更快)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 3. 训练循环(10轮)
epochs = 10
for t in range(epochs):
    print(f"\nEpoch {t+1}\n-------------------------------")
    trainda(train_dataloader, model, loss_fn, optimizer)
    # testda(test_dataloader, model, loss_fn)  # 可选:每轮训练后评估
print("Training Done!")

# 最终评估
testda(test_dataloader, model, loss_fn)
关键参数说明:
  • 学习率(lr):0.001 是 Adam 优化器的经典初始值,过大易震荡,过小收敛慢;
  • 批次大小(batch_size):64 兼顾显存占用与训练稳定性,可根据 GPU 显存调整;
  • 训练轮数(epochs):10 轮为基础值,可根据验证集损失调整(避免过拟合)。

三、常见问题与优化建议

核心问题排查

  • 维度不匹配:全连接层输入维度需与卷积层输出匹配(如 500×350 输入需改为128x125x87);
  • 标签类型错误:CrossEntropyLoss 要求标签为 64 位整数,需用dtype = np.int64;
  • 设备不匹配:确保数据(X/Y)与模型在同一设备(GPU/CPU);
  • 图像读取失败:添加异常处理(try-except),跳过损坏的图像文件。

四、总结

本文以食品图像分类为例,完整实现了基于 PyTorch 的 CNN 训练流程,核心要点包括:

  1. 自定义 Dataset 需实现三大核心方法,确保数据正确加载与预处理;
  2. CNN 设计需遵循「卷积提取特征→池化降维→全连接分类」的逻辑,参数设计需匹配输入尺寸;
  3. 训练过程需区分训练 / 评估模式,关闭评估阶段的梯度计算以节省资源;
  4. 实际应用中需关注维度匹配、设备一致性、数据增强等关键细节,提升模型性能。
相关推荐
lijianhua_97126 小时前
国内某顶级大学内部用的ai自动生成论文的提示词
人工智能
EDPJ6 小时前
当图像与文本 “各说各话” —— CLIP 中的模态鸿沟与对象偏向
深度学习·计算机视觉
蔡俊锋6 小时前
用AI实现乐高式大型可插拔系统的技术方案
人工智能·ai工程·ai原子能力·ai乐高工程
自然语6 小时前
人工智能之数字生命 认知架构白皮书 第7章
人工智能·架构
大熊背7 小时前
利用ISP离线模式进行分块LSC校正的方法
人工智能·算法·机器学习
eastyuxiao7 小时前
如何在不同的机器上运行多个OpenClaw实例?
人工智能·git·架构·github·php
诸葛务农7 小时前
AGI 主要技术路径及核心技术:归一融合及未来之路5
大数据·人工智能
光影少年7 小时前
AI Agent智能体开发
人工智能·aigc·ai编程
ai生成式引擎优化技术7 小时前
TSPR-WEB-LLM-HIC (TWLH四元结构)AI生成式引擎(GEO)技术白皮书
人工智能
帐篷Li7 小时前
9Router:开源AI路由网关的架构设计与技术实现深度解析
人工智能