基于 PyTorch 的食物图像分类CNN 训练全流程

本文将从零开始带大家实现一个基于 PyTorch 的 20 类食物图像分类项目，涵盖数据集准备→自定义 Dataset→CNN 网络搭建→模型训练→测试评估完整流程。

我们将按照工业界标准流程，一步步完成：

数据集标签文件自动生成
PyTorch 自定义 Dataset 类构建
DataLoader 批量数据加载
CNN 卷积神经网络搭建
模型训练与测试评估

全程附完整可运行代码，以及 Windows 环境下常见坑点解决方案。

一、数据集标签文件生成

二、依赖库导入与图片预处理

[三、自定义数据集类 food_dataset](#三、自定义数据集类 food_dataset)

[四、DataLoader 数据加载与 CNN 网络搭建](#四、DataLoader 数据加载与 CNN 网络搭建)

五、模型训练与测试完整流程

一、数据集标签文件生成

在深度学习训练中，我们需要建立「图片路径 → 分类标签」的映射关系。最通用、最易维护的方式就是生成标签 txt 文件，每行格式如下图：图片完整路径＋数字标签

本项目要求数据集遵循以下层级结构：

train 和 test 目录下必须是一级类别子文件夹，图片直接放在类别文件夹内。禁止直接在 train 下放图片，也禁止多层嵌套，否则后面的遍历脚本会乱。

接下来就是自动生成标签文件的代码。我写了一个函数叫train_test_file，给它传入数据集根目录和需要处理的子集名（'train'，'test'），它就会在同级目录下生成对应的train.txt或test.txt。

python 复制代码

import os

def train_test_file(root, dir):
    """
    自动遍历数据集文件夹，生成图片路径+数字标签的txt文件
    :param root: 数据集根目录路径
    :param dir: 子集名称，传入'train'或'test'
    """
    # 创建标签文件，强制指定utf-8编码，解决Windows中文路径乱码
    file_txt = open(dir + '.txt', 'w', encoding="utf-8")
    # 拼接数据集完整路径
    path = os.path.join(root, dir)

    # os.walk递归遍历文件夹，返回(当前路径, 子文件夹列表, 文件列表)
    for roots, directories, files in os.walk(path):
        # 分支1：首次进入train/test目录，获取所有类别名称
        if len(directories) != 0:
            dirs = directories
        # 分支2：进入具体类别文件夹，开始遍历图片写入标签
        else:
            # 分割路径提取当前类别名称
            now_dir = roots.split('\\')
            for file in files:
                # 拼接图片完整路径
                path_1 = os.path.join(roots, file)
                print(path_1)
                # 写入格式：图片路径 类别下标（下标作为数字标签）
                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)

中文乱码解决方案

复制代码

file_txt = open(dir + '.txt', 'w', encoding="utf-8")

如果你在Windows下不写encoding="utf-8"，默认会用系统本地编码（比如GBK），而我们的文件夹名是中文，写入txt时就会变成乱码，后面用Dataset读取路径时直接报错"文件不存在"。

os.walk 双分支设计

复制代码

if len(directories) != 0:
    dirs = directories  # 保存所有类别名称
else:
    # 进入图片文件夹，写入标签

利用 os.walk 先遍历到顶层目录获取全部类别，再进入子目录生成标签，os.walk每次返回三个值：当前目录路径roots、当前目录下的子文件夹列表directories、当前目录下的文件列表files。我们利用它的两层特性：第一次进入train目录时，directories不为空，里面正好是20个类别文件夹的名字，我们把这20个名字存到dirs变量里。然后os.walk会继续往下走，进入每一个类别文件夹，这时directories为空，就走到了else分支，开始逐个处理图片文件。

标签生成原理

复制代码

dirs.index(now_dir[-1])

用类别名称在类别列表中的下标作为数字标签，保证相同类别标签一致。now_dir是通过roots.split('\\')得到的路径碎片列表，最后一项now_dir[-1]就是当前所在的类别文件夹名，比如"汉堡"。dirs是之前保存的所有类别名的列表，顺序就是os.walk遍历到的顺序。用index方法找到"汉堡"在dirs中的位置，这个位置数字就是我们要的数字标签。

运行完脚本后，会在当前目录生成train.txt和test.txt，控制台也会逐行打印所有找到的图片路径，如果你发现txt文件是空的，先检查root路径对不对，再看看train文件夹下面是否确实有类别子文件夹，而不是直接把图片扔在train下面。另外，路径分隔符在Windows下是反斜杠，代码里直接写'\\'或者用os.path.join自动处理都可以。

二、依赖库导入与图片预处理

python 复制代码

import torch
from torch.utils.data import Dataset, DataLoader  # 数据集基类 + 批次加载器
import numpy as np
from PIL import Image                              # 图片读取库
from torchvision import transforms                 # 图片预处理工具

定义图片预处理

transforms.Compose( $...$ )，你可以把它想象成一条流水线，把你想做的所有预处理按顺序写进一个列表里，图片从这头进去，就会依次经过每个操作，从那头出来就是模型能直接接受的张量了。

python 复制代码

data_transforms = {
    # 训练集预处理流水线
    'trainda':
        transforms.Compose([
            transforms.Resize([256, 256]),  # 统一缩放为256×256像素
            transforms.ToTensor(),          # PIL图片 → Tensor张量
        ]),
    # 测试集预处理流水线
    'valid':
        transforms.Compose([
            transforms.Resize([256, 256]),
            transforms.ToTensor(),
        ]),
}

transforms.Resize 尺寸统一。原始图片的分辨率五花八门，有的可能很大，有的很小，但神经网络的全连接层要求输入尺寸必须完全一致，所以我们要把所有图片强制缩放到同一个大小。这里选的是 $256, 256$ ，也就是宽高各256像素。为什么是256呢？因为像 VGG、ResNet 这些经典 CNN 网络的标准输入尺寸往往是 224×224 或 256×256，256 算是一个兼顾了计算精度和处理速度的选择。当然你也可以改成其他尺寸，只要和网络定义的第一层匹配就行。

**transforms.ToTensor 格式转换，**这一步做了三件大事，而且缺一不可。

第一，数据类型转换：PIL 库读取的图片对象会被转成 PyTorch 的张量（Tensor）。

第二，数值归一化：原始图片每个像素的取值是 0 到 255 的整数，ToTensor 会自动除以 255，把它们变成 0.0 到 1.0 的浮点数。深度学习模型对输入数值范围很敏感，如果输入是 0~255 的大数，梯度更新很容易不稳定。

第三，维度重排。PIL 图片的维度顺序是 (高度, 宽度, 通道)，比如一张 256×256 的彩色图片，它的形状是 (256, 256, 3)，三个通道（RGB）在最后。但 PyTorch 里的卷积层 nn.Conv2d 要求输入张量的形状是 (通道, 高度, 宽度)，也就是把通道数挪到最前面。ToTensor 会帮你自动完成这个重排，输出形状变成 (3, 256, 256)。之所以要通道前置，是因为 PyTorch 的设计惯例是批量维度在第一维，通道维度在第二维，后面才是空间维度。这种做法在底层内存访问上更高效，也方便卷积核跨通道计算。

三、自定义数据集类 food_dataset

在 PyTorch 中，自定义数据集是深度学习项目的核心环节。通过继承官方的 Dataset 基类，我们可以实现灵活、高效的数据读取逻辑。

PyTorch 提供的 Dataset 是一个抽象基类，它定义了数据集的标准接口，任何自定义 Dataset 都必须实现以下三个方法：

python 复制代码

class food_dataset(Dataset):
    """
    继承PyTorch官方Dataset基类，实现自定义数据集读取
    必须实现3个方法：__init__, __len__, __getitem__
    """

    def __init__(self, file_path, transform=None):
        """
        初始化：读取标签txt文件，保存所有图片路径和对应标签
        """
        self.file_path = file_path    # 标签文件路径
        self.imgs = []                # 存储所有图片路径
        self.labels = []              # 存储所有对应标签
        self.transform = transform    # 图片预处理流水线

        # 读取标签txt文件，每行按空格分割为 [图片路径, 数字标签]
        with open(self.file_path, encoding="utf-8") 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):
        """按索引返回单条样本：(预处理后的图片张量, 标签张量)"""
        # 1. 根据索引读取图片
        image = Image.open(self.imgs[idx])
        # 2. 如果传入了预处理流水线，执行图片变换
        if self.transform:
            image = self.transform(image)
        # 3. 标签转换：字符串 → int64张量（分类任务标准格式）
        label = self.labels[idx]
        label = torch.from_numpy(np.array(label, dtype=np.int64))
        return image, label

实现一个自定义的 Dataset 类。这个类需要定义三个核心方法，分别是初始化方法、长度方法和取样本方法。

init 初始化方法

这里的 file_path 参数传入 train.txt 或 test.txt 的路径，transform 参数则传入上一节定义的预处理流水线。

文件读取的逻辑是这样的。先用 f.readlines () 读取 txt 中的所有行，然后对每一行用 strip () 去掉首尾的换行和空格，再用 split (' ') 按空格分割，最终得到图片路径和标签的列表。

len 长度方法

这个方法最简单，直接返回图片总数即可。当你执行 len (training_data) 时，PyTorch 会自动调用这个方法。

getitem 取样本方法

这是 Dataset 最核心的方法，每次按索引取数据时都会执行。

代码中的第一步是根据索引 idx 取出对应的图片路径，用 Image.open 读取得到 PIL 格式的图片。如果传入了 transform 预处理流水线，就调用它依次完成缩放、转 Tensor 和通道重排。

接着处理标签。从 txt 文件中读出的标签是字符串 "0"、"1" 这样的形式，需要先转换成 numpy 的 int64 数组，再转换成 PyTorch 张量。之所以要用 int64，是因为分类任务中常用的交叉熵损失函数要求标签必须是 long 类型也就是 int64，这是标准的输入格式。

init 方法中只保存图片路径，并不读取图片，直到 getitem 被调用时才真正将图片读入内存。这样一来，即使数据集有几万张图片，内存占用也会非常小。

四、DataLoader 数据加载与 CNN 网络搭建

实例化数据集与 DataLoader

python 复制代码

# 实例化训练集、测试集
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)

首先实例化训练集和测试集，分别传入对应的标签文件路径和预处理流水线。

然后通过 DataLoader 将数据集打包成批次。batch_size=64 表示每次训练一次性送入 64 张图片。shuffle=True 表示训练集打乱顺序，防止模型记住图片排列顺序。

这里需要注意测试集通常设置 shuffle=False，保证每次评估顺序固定，准确率结果稳定。

自动选择训练设备

python 复制代码

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

这段代码会自动判断当前硬件环境，优先使用 NVIDIA 显卡的 cuda 加速，其次是苹果 M 系列芯片的 mps 加速，最后兜底使用 CPU。后续所有数据和模型都会迁移到这个设备上运算。

搭建 CNN 卷积神经网络

python 复制代码

from torch import nn

class CNN(nn.Module):
    """自定义CNN卷积神经网络，20分类"""

    def __init__(self):
        super(CNN, self).__init__()
        # 第1组卷积块：卷积+激活+池化
        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)
        )
        # 第2组卷积块：两次卷积+激活+池化
        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组卷积块：仅卷积+激活
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=128, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
        )
        # 全连接输出层：映射到20个分类
        self.out = nn.Linear(128 * 64 * 64, out_features=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

# 实例化模型并迁移到指定设备
model = CNN().to(device)
print(model)

网络结构采用三组卷积块逐步提取图像特征。

第一组卷积将 3 通道 RGB 图映射为 16 张特征图，池化后尺寸从 256 缩小到 128。第二组两次卷积通道数提升到 32，再次池化尺寸缩小到 64。第三组卷积继续升维到 128 通道，保持尺寸不变。

最后通过 view 将四维特征图展平，送入全连接层输出 20 个类别的预测分数。padding=2 保证卷积后尺寸不变，这是卷积网络的常用技巧。

五、模型训练与测试完整流程

单轮训练函数

python 复制代码

def trainda(dataloader, model, loss_fn, optimizer):
    """
    功能：执行单轮完整训练
    流程：前向传播 → 计算损失 → 反向传播 → 更新权重
    """
    model.train()
    batch_size_num = 1

    for X, y in dataloader:
        # 数据迁移到GPU/CPU
        X, y = X.to(device), y.to(device)
        # 前向传播得到预测结果
        pred = model.forward(X)
        # 计算预测与真实标签的损失
        loss = loss_fn(pred, y)

        # 反向传播更新权重标准三步
        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

训练函数是深度学习的核心。首先调用 model.train () 切换到训练模式。然后循环读取每个批次的数据，将数据和标签迁移到加速设备。

前向传播得到预测结果，计算预测与真实标签的损失。反向传播有标准的三步：先清空上一轮梯度缓存，再反向传播计算梯度，最后优化器根据梯度更新权重。这三步顺序不能错。

测试集评估函数

python 复制代码

def testda(dataloader, model, loss_fn):
    """
    功能：在测试集上评估模型准确率和平均损失
    """
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0

    with torch.no_grad():
        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()

    test_loss /= num_batches
    correct /= size
    print(f"Test result: \n Accuracy: {(100 * correct)}%, Avg loss: {test_loss}")

测试函数首先调用 model.eval () 切换到评估模式，关闭 Dropout 等训练时的随机行为。

最重要的是 torch.no_grad ()，测试不需要更新权重，关闭梯度计算可以大幅节省显存并提升推理速度。pred.argmax (1) 取每行最大值的下标作为预测类别，与真实标签对比统计正确数量。

定义损失函数与优化器

python 复制代码

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

多分类任务标准使用 CrossEntropyLoss 交叉熵损失，内部自动完成 softmax 归一化。优化器选择 Adam，自带自适应学习率，收敛速度快，是图像分类的首选。

完整训练循环

python 复制代码

epochs = 10
for t in range(epochs):
    print(f"Epoch {t + 1}\n--------------------------------")
    trainda(train_dataloader, model, loss_fn, optimizer)
    testda(test_dataloader, model, loss_fn)

print("Done!")
testda(test_dataloader, model, loss_fn)

一个 epoch 代表把整个训练集完整遍历一遍。这里设置训练 10 轮，每轮训练完成后立即在测试集上评估，这样可以实时观察模型是否过拟合。全部训练结束后再执行一次最终评估。