深度学习数据加载实战：从 PyTorch Dataset 到食品图像分类全流程解析

一、开篇：为什么数据加载是深度学习的 “地基”？

在训练一个图像分类模型时，我们通常需要完成三个核心步骤：数据准备→模型训练→评估优化。其中，数据加载是连接 “原始数据” 与 “模型输入” 的桥梁 —— 如果桥梁不稳（比如数据格式错误、批量加载卡顿），再复杂的模型也无法发挥作用。

举个真实场景：某同学在训练食品分类模型时，因未统一图像尺寸，导致输入张量形状不一致，模型直接报错；另一同学因未使用多进程加载（num_workers=0），在 10 万张图像的数据集上，单轮 epoch 加载时间长达 20 分钟，训练效率极低。

这些问题的根源，在于对 PyTorch 数据加载机制的理解不深入。PyTorch 为我们提供了一套优雅的解决方案：

• Dataset：定义 “数据从哪里来、如何读取单个样本”；
• DataLoader：定义 “如何批量加载样本、是否打乱、是否多进程加速”；
• torchvision.transforms：定义 “图像如何预处理（如 Resize、归一化）”。

本文将围绕食品图像分类任务，从这三个核心组件出发，逐步拆解数据加载的全流程，最终实现一个可直接复用的工业级数据处理方案。

二、基础原理：PyTorch 数据加载的 “三驾马车”

在解析代码前，我们先理清 PyTorch 数据加载的核心逻辑。简单来说，数据加载的流程是：原始数据→Dataset（单样本处理）→DataLoader（批量处理）→模型输入。下面逐一拆解每个组件的作用。

2.1 Dataset：定义 “单样本的读取规则”

Dataset是 PyTorch 中所有自定义数据集的基类，它规定了数据集必须实现两个核心方法：__len__和__getitem__，同时可通过__init__初始化参数。

2.1.1 Dataset 的核心方法解析

• __init__：初始化数据集，通常用于读取标签文件、定义数据路径、绑定预处理函数（transform）。比如在食品数据集中，我们需要通过__init__读取存储 “图像路径 - 标签” 的 txt 文件，将路径和标签分别存入列表。
• __len__：返回数据集的总样本数，用于告诉 DataLoader “数据集有多大”，从而计算批次数量。
• __getitem__：根据索引idx返回单个样本（图像张量 + 标签），是 Dataset 的核心。它需要完成：读取图像→应用预处理→处理标签→返回样本。

为什么必须实现这两个方法？因为 DataLoader 在迭代加载数据时，会通过__len__获取总样本数，再通过__getitem__(idx)逐个读取样本，最后组装成批次。如果缺少任一方法，DataLoader 将无法正常工作。

2.1.2 内置 Dataset vs 自定义 Dataset

PyTorch 提供了一些内置 Dataset（如ImageFolder、CIFAR10），其中ImageFolder适用于 “按类别划分文件夹” 的数据集（如food/train/汉堡/001.jpg、food/train/披萨/002.jpg）。但实际项目中，数据格式往往更灵活（如用 txt 文件记录路径和标签），此时就需要自定义 Dataset—— 这也是本文代码的核心场景。

2.2 DataLoader：实现 “批量加载与效率优化”

Dataset仅负责单样本的读取，而DataLoader则是在Dataset的基础上，实现批量加载、数据打乱、多进程加速、内存优化等高级功能。它相当于一个 “数据搬运工”，将Dataset生产的单样本，打包成模型需要的批次数据。

2.2.1 DataLoader 的关键参数详解

在本文代码中，我们用DataLoader封装自定义数据集时，用到了几个核心参数：

python

train_loader = DataLoader(train_dataset,    # 传入自定义Dataset实例batch_size=32,    # 每批次样本数shuffle=True,     # 训练集是否打乱样本顺序num_workers=4     # 用于加载数据的进程数
)

这些参数直接影响训练效率和模型性能，必须深入理解：

参数	作用	注意事项
dataset	传入的 Dataset 实例，是数据的来源	必须实现__len__和__getitem__方法
batch_size	每批次加载的样本数	需根据 GPU 内存调整（如 12GB 内存可设 32/64，8GB 内存设 16/32），过大会导致 OOM
shuffle	是否在每个 epoch 前打乱样本顺序	训练集设True（避免模型记忆样本顺序），验证集设False（保证结果可复现）
num_workers	用于加载数据的进程数（默认为 0，即主进程加载）	设为 CPU 核心数的 1~2 倍（如 4 核 CPU 设 4），Windows 下需注意多进程兼容性
pin_memory	是否将加载的数据存入 GPU 锁页内存（默认为 False）	使用 GPU 时设True，可加速数据从 CPU 到 GPU 的传输（减少内存拷贝时间）
drop_last	是否丢弃最后一个不足batch_size的批次（默认为 False）	训练集可设True（避免批次大小不一致影响 BN 层），验证集设False（不浪费数据）

2.2.2 DataLoader 的工作流程

1. 每个 epoch 开始时，若shuffle=True，打乱样本索引；
2. 根据batch_size将索引分成多个批次；
3. 启动num_workers个进程，并行读取每个批次的样本（通过__getitem__）；
4. 将读取的样本组装成批次张量（形状为[batch_size, channels, height, width]）；
5. 若pin_memory=True，将张量存入 GPU 锁页内存，等待模型调用。

2.3 torchvision.transforms：图像预处理的 “流水线”

原始图像（如 JPG/PNG 文件）无法直接输入模型 —— 它们的尺寸可能不一致（如 200×300、500×400）、像素值范围为[0,255]（而模型需要[0,1]或归一化后的数值）、格式为 PIL 图像（模型需要张量）。torchvision.transforms就是用于解决这些问题的 “预处理流水线”。

2.3.1 常用 transforms 操作解析

本文代码中用到了Resize和ToTensor，但实际项目中需要更完整的预处理流程。以下是食品图像分类中常用的 transforms 操作：

操作	作用	代码示例
Resize(size)	将图像缩放到指定尺寸（如[256,256]）	transforms.Resize([256,256])
RandomResizedCrop	随机裁剪到指定尺寸（训练集数据增强，增加多样性）	transforms.RandomResizedCrop(224)
RandomHorizontalFlip	随机水平翻转（概率 0.5，数据增强）	transforms.RandomHorizontalFlip(p=0.5)
ToTensor()	将 PIL 图像（H×W×C）转为 PyTorch 张量（C×H×W），并将像素值归一化到[0,1]	transforms.ToTensor()
Normalize(mean, std)	对张量进行归一化（output = (input - mean) / std）	transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225])
ColorJitter	随机调整亮度、对比度、饱和度（数据增强）	transforms.ColorJitter(brightness=0.2, contrast=0.2)

2.3.2 训练集 vs 验证集：预处理的差异

训练集需要数据增强（如随机裁剪、翻转、颜色抖动），目的是增加数据多样性，防止模型过拟合；验证集则不需要数据增强，只需统一尺寸、归一化，确保评估结果的客观性。

例如，在本文代码的基础上，我们可以优化预处理流程：

python

data_transforms = {'train': transforms.Compose([transforms.RandomResizedCrop(224),  # 随机裁剪到224×224（适配预训练模型）transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 颜色抖动transforms.ToTensor(),  # 转张量并归一化到[0,1]transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 标准化]),'valid': transforms.Compose([transforms.Resize([256, 256]),  # 先缩放到256×256transforms.CenterCrop(224),  # 中心裁剪到224×224（与训练集一致）transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
}

这里用到的Normalize参数（均值[0.485,0.456,0.406]、标准差[0.229,0.224,0.225]）来自 ImageNet 数据集 —— 如果使用预训练模型（如 ResNet18），必须使用与预训练数据一致的归一化参数，否则模型性能会大幅下降。

三、深度解析：食品图像数据集代码的每一行

在理解了核心原理后，我们来逐行解析本文的核心代码 —— 自定义食品图像数据集（food_dataset）。代码看似简短，但每一行都暗藏细节，新手容易在这里踩坑。

3.1 代码整体结构回顾

首先，我们先看完整的带注释代码（已优化预处理流程）：

python

# 1. 导入依赖库
import torch  # PyTorch核心库（张量操作、模型构建）
from torch.utils.data import Dataset, DataLoader  # 数据加载核心组件
import numpy as np  # 数值计算库（用于标签格式转换）
from PIL import Image  # 图像读取库（处理JPG/PNG文件）
from torchvision import transforms  # 图像预处理库
import os  # 文件路径处理库（新增，用于路径验证）# 2. 定义图像预处理流水线
data_transforms = {'train': transforms.Compose([transforms.RandomResizedCrop(224),  # 训练集：随机裁剪到224×224（数据增强）transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转（概率0.5）transforms.ColorJitter(brightness=0.2, contrast=0.2),  # 随机调整亮度、对比度transforms.ToTensor(),  # 转换为张量：H×W×C → C×H×W，像素值[0,255]→[0,1]# 标准化：使用ImageNet均值方差（适配预训练模型）transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),'valid': transforms.Compose([transforms.Resize([256, 256]),  # 验证集：先缩放到256×256transforms.CenterCrop(224),  # 中心裁剪到224×224（与训练集输入尺寸一致）transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
}# 3. 自定义食品图像数据集类（继承Dataset基类）
class food_dataset(Dataset):"""食品图像数据集类，用于加载食品图像及其对应的分类标签适用场景：标签存储在txt文件中，每行格式为"图像路径 标签"（如"food/001.jpg 0"）核心功能：读取图像、应用预处理、返回张量格式的样本"""def __init__(self, file_path, transform=None, check_path=True):"""初始化数据集参数：file_path (str): 标签txt文件的路径（如"train_labels.txt"）transform (callable, optional): 图像预处理函数（默认None，即不预处理）check_path (bool, optional): 是否验证图像路径有效性（默认True，避免路径错误）"""# 保存参数到实例属性self.file_path = file_pathself.transform = transformself.check_path = check_path# 初始化存储图像路径和标签的列表self.imgs = []  # 存储所有图像的绝对/相对路径self.labels = []  # 存储所有图像对应的整数标签# 4. 读取标签文件并解析路径和标签try:# 打开txt文件（使用with语句自动关闭文件，避免资源泄漏）with open(self.file_path, 'r', encoding='utf-8') as f:# 读取所有行，去除首尾空白字符（如换行符\n、空格）lines = [line.strip() for line in f.readlines()]# 过滤空行（避免因txt文件末尾空行导致错误）lines = [line for line in lines if line]# 遍历每一行，解析图像路径和标签for line in lines:# 按空格分割路径和标签（假设每行只有一个空格分隔）parts = line.split(' ')# 异常处理：若行格式不正确（如无空格、多空格），跳过并提示if len(parts) != 2:print(f"警告：跳过格式错误的行 -> {line}")continueimg_path, label_str = parts# 验证图像路径是否存在（可选，但推荐开启）if self.check_path and not os.path.exists(img_path):print(f"警告：图像路径不存在 -> {img_path}，跳过该样本")continue# 将路径添加到列表self.imgs.append(img_path)# 将标签从字符串转为整数（分类任务标签需为int64类型）try:label = int(label_str)self.labels.append(label)except ValueError:print(f"警告：标签不是整数 -> {label_str}，跳过该样本")# 若标签无效，需删除对应的路径（保持列表长度一致）self.imgs.pop()except FileNotFoundError:# 若标签文件不存在，抛出异常并提示raise FileNotFoundError(f"标签文件不存在：{self.file_path}")except Exception as e:# 捕获其他异常（如权限错误）raise RuntimeError(f"读取标签文件失败：{str(e)}")# 验证路径和标签列表长度是否一致（避免后续索引不匹配）assert len(self.imgs) == len(self.labels), \f"图像路径列表长度（{len(self.imgs)}）与标签列表长度（{len(self.labels)}）不匹配"# 打印数据集初始化信息（方便调试）print(f"数据集初始化完成：共加载 {len(self.imgs)} 个样本")def __len__(self):"""返回数据集的总样本数（必须实现）返回：int: 样本总数"""return len(self.imgs)def __getitem__(self, idx):"""根据索引idx返回单个样本（必须实现）参数：idx (int): 样本索引（0 <= idx < __len__()）返回：tuple: (image_tensor, label_tensor)，即预处理后的图像张量和标签张量"""# 1. 根据索引获取图像路径和标签img_path = self.imgs[idx]label = self.labels[idx]# 2. 读取图像（处理可能的读取错误）try:# 用PIL打开图像（默认打开为RGB格式，若为灰度图会自动转为单通道）image = Image.open(img_path).convert('RGB')  # 强制转为RGB（避免灰度图通道数不一致）except Exception as e:# 若图像读取失败，抛出异常并提示路径raise RuntimeError(f"读取图像失败：{img_path}，错误信息：{str(e)}")# 3. 应用预处理（若有）if self.transform is not None:try:image = self.transform(image)except Exception as e:raise RuntimeError(f"图像预处理失败：{img_path}，错误信息：{str(e)}")# 4. 处理标签（转换为int64类型张量）# 为什么用int64？因为PyTorch的CrossEntropyLoss要求标签为torch.int64类型label_tensor = torch.tensor(label, dtype=torch.int64)# 5. 返回样本（图像张量+标签张量）return image, label_tensor# 4. 测试数据集和数据加载器（仅在当前脚本运行时执行）
if __name__ == "__main__":# 4.1 配置路径（根据实际项目修改）train_label_path = "train_labels.txt"  # 训练集标签文件路径valid_label_path = "valid_labels.txt"  # 验证集标签文件路径# 4.2 创建训练集和验证集实例print("=== 初始化训练集 ===")train_dataset = food_dataset(file_path=train_label_path,transform=data_transforms['train'],check_path=True  # 开启路径验证)print("\n=== 初始化验证集 ===")valid_dataset = food_dataset(file_path=valid_label_path,transform=data_transforms['valid'],check_path=True)# 4.3 创建数据加载器train_loader = DataLoader(train_dataset,batch_size=32,  # 每批次32个样本shuffle=True,   # 训练集打乱num_workers=4,  # 4个进程加载数据pin_memory=True,  # 开启锁页内存（若使用GPU）drop_last=True  # 丢弃最后一个不足32的批次)valid_loader = DataLoader(valid_dataset,batch_size=32,shuffle=False,  # 验证集不打乱num_workers=4,pin_memory=True,drop_last=False  # 不丢弃最后一个批次（避免浪费数据）)# 4.4 测试数据加载器（迭代一个批次）print("\n=== 测试训练集数据加载器 ===")for batch_idx, (images, labels) in enumerate(train_loader):# 打印批次信息print(f"批次索引：{batch_idx}")print(f"图像张量形状：{images.shape}")  # 输出：torch.Size([32, 3, 224, 224])print(f"标签张量形状：{labels.shape}")  # 输出：torch.Size([32])print(f"标签取值范围：{labels.min()} ~ {labels.max()}")  # 验证标签是否正确# 仅测试一个批次（避免耗时）breakprint("\n=== 测试验证集数据加载器 ===")for batch_idx, (images, labels) in enumerate(valid_loader):print(f"批次索引：{batch_idx}")print(f"图像张量形状：{images.shape}")print(f"标签张量形状：{labels.shape}")break# 4.5 验证GPU兼容性（若有GPU）if torch.cuda.is_available():device = torch.device("cuda:0")# 将一个批次的数据转移到GPUimages = images.to(device)labels = labels.to(device)print(f"\n=== GPU测试 ===")print(f"图像张量设备：{images.device}")  # 输出：cuda:0print(f"标签张量设备：{labels.device}")  # 输出：cuda:0print("GPU数据转移成功！")

3.2 关键代码细节解析（新手必看）

3.2.1 依赖库导入：为什么需要这些库？

• torch：所有 PyTorch 功能的基础，用于创建张量、定义模型等；
• Dataset/DataLoader：数据加载的核心组件，必须从torch.utils.data导入（新手常漏导Dataset，导致NameError）；
• numpy：本文中用于标签转换（虽然后续优化为直接用torch.tensor，但 numpy 在处理大规模数据时更灵活）；
• PIL.Image：Python 中最常用的图像读取库，支持多种格式（JPG、PNG、BMP 等）；
• os：新增的路径处理库，用于验证图像路径是否存在（避免因路径错误导致的FileNotFoundError）。

3.2.2 预处理流水线：为什么要这么设计？

1. RandomResizedCrop(224)：
   训练集使用随机裁剪，而非固定Resize，是因为随机裁剪可以让模型看到图像的不同区域（如食品的局部特征），增强模型的泛化能力。例如，一张汉堡图像可能被裁剪到面包部分，也可能被裁剪到肉饼部分，模型不会依赖固定的图像区域。

2. convert('RGB')：
   强制将图像转为 RGB 三通道格式，避免灰度图（单通道）导致的通道数不一致。例如，某些食品图像可能是灰度图（如老照片），若不转为 RGB，其张量形状会是[1,224,224]，而 RGB 图像是[3,224,224]，混合后会导致批次张量形状错误。

3. Normalize的数学原理：
   归一化的公式是output = (input - mean) / std，目的是将像素值从[0,1]转为零均值、单位方差的分布。例如，对于 RGB 图像的 R 通道，原始值为 0.5，均值为 0.485，标准差为 0.229，则归一化后的值为(0.5 - 0.485)/0.229 ≈ 0.0655。这种处理能加速模型收敛，因为模型对数值范围敏感，过大或过小的数值会导致梯度爆炸 / 消失。

3.2.3 自定义 Dataset 的异常处理：为什么要加这么多 try-except？

新手写 Dataset 时，常忽略异常处理，导致程序在遇到错误时直接崩溃，且难以定位问题。本文代码中添加了多层异常处理：

1. 标签文件读取异常：若file_path不存在，直接抛出FileNotFoundError并提示路径，避免用户找不到错误原因；
2. 行格式错误：若 txt 文件中有行格式不正确（如无空格、多空格），跳过该行并提示，避免整个数据集加载失败；
3. 图像路径无效：通过os.path.exists(img_path)验证路径，跳过无效路径，避免后续读取图像时崩溃；
4. 标签类型错误：若标签不是整数（如字符串 “汉堡”），跳过该样本并提示，避免int(label_str)报错；
5. 图像读取失败：若图像损坏或格式不支持，抛出RuntimeError并提示路径，方便用户排查损坏文件。

这些异常处理让数据集加载更稳健，尤其在处理大规模数据（如 10 万 + 样本）时，能避免因个别坏样本导致整个训练中断。

3.2.4 标签处理：为什么要用torch.int64？

PyTorch 的CrossEntropyLoss（交叉熵损失）要求标签必须是torch.int64类型（即长整型），不能是int32或float。若使用其他类型，会抛出RuntimeError: expected scalar type Long but found Int。

本文中用torch.tensor(label, dtype=torch.int64)直接创建 int64 类型的标签张量，替代了原代码中的torch.from_numpy(np.array(label,dtype=np.int64))，两种方式效果一致，但前者更简洁。

3.2.5 if __name__ == "__main__":的作用

这是 Python 的常用语法，用于判断脚本是否被直接运行（而非被导入为模块）。在该代码块中，我们测试数据集和数据加载器的功能，验证：

• 数据集是否能正常初始化（样本数量是否正确）；
• 数据加载器是否能正常批量加载（张量形状是否正确）；
• 数据是否能正常转移到 GPU（若有 GPU）。

这种测试能帮助我们在训练模型前，提前发现数据加载中的问题，避免在训练过程中（尤其是训练多个 epoch 后）才报错。

四、实战优化：让数据加载更快、更稳健

在实际项目中，仅实现基础的数据集和数据加载器还不够 —— 当数据集规模达到 10 万 + 样本时，加载速度会成为瓶颈；当数据存在不均衡（如某些食品类别样本极少）时，模型会偏向多数类。本节将介绍 5 个实战优化技巧，让数据加载流水线更高效、更稳健。

4.1 优化 1：多进程加载与内存锁页（提升加载速度）

num_workers和pin_memory是影响加载速度的关键参数，正确设置能大幅减少数据加载时间。

4.1.1 num_workers的最佳实践

• 设置原则：num_workers的取值通常为 CPU 核心数的 1~2 倍。例如，4 核 CPU 设 4，8 核 CPU 设 8 或 16。
• Windows 注意事项：Windows 系统下，多进程加载可能会出现BrokenPipeError，此时可尝试：
  1. 将num_workers设为 0（仅主进程加载，速度慢但稳定）；
  2. 在脚本开头添加if __name__ == "__main__":（避免多进程重复初始化）；
  3. 使用 PyTorch 1.8 + 版本（修复了 Windows 多进程的部分 bug）。

4.1.2 pin_memory=True的作用

当使用 GPU 训练时，pin_memory=True会将加载的数据存入 GPU 的锁页内存（page-locked memory），而非普通内存。普通内存的数据转移到 GPU 时，需要先复制到锁页内存，再转移到 GPU；而锁页内存的数据可以直接转移到 GPU，减少一次内存拷贝，从而提升速度。

注意：若使用 CPU 训练，pin_memory=True无效，且会占用更多内存，建议设为False。

4.2 优化 2：数据预加载与缓存（减少重复读取）

当数据集规模较大时，每次 epoch 都从硬盘读取图像会耗时较长。我们可以将预处理后的图像缓存到内存或固态硬盘（SSD）中，减少重复读取。

4.2.1 内存缓存（适用于小数据集）

在__init__中提前读取所有图像并应用预处理，存入内存：

python

def __init__(self, file_path, transform=None, check_path=True):# ... 其他初始化代码 ...# 内存缓存：提前读取所有图像并预处理self.cache = []for img_path, label in zip(self.imgs, self.labels):image = Image.open(img_path).convert('RGB')if transform is not None:image = transform(image)self.cache.append( (image, torch.tensor(label, dtype=torch.int64)) )def __getitem__(self, idx):# 直接从缓存中获取样本，无需重复读取和预处理return self.cache[idx]

优点：加载速度极快，适合小数据集（如 1 万样本以内）；
缺点：占用大量内存（每张 224×224 的 RGB 图像约 224×224×3=150KB，10 万样本约 15GB）。

4.2.2 磁盘缓存（适用于大数据集）

将预处理后的图像保存为 PyTorch 的.pt格式（张量文件），下次加载时直接读取张量：

python

def __init__(self, file_path, transform=None, cache_dir="cache", check_path=True):# ... 其他初始化代码 ...# 创建缓存目录os.makedirs(cache_dir, exist_ok=True)self.imgs = []self.labels = []with open(file_path, 'r') as f:lines = [line.strip() for line in f.readlines()]for line in lines:img_path, label_str = line.split(' ')label = int(label_str)# 生成缓存文件名（用图像路径的哈希值避免重复）img_hash = hashlib.md5(img_path.encode()).hexdigest()cache_path = os.path.join(cache_dir, f"{img_hash}.pt")# 若缓存存在，直接读取；否则预处理后保存if os.path.exists(cache_path):image = torch.load(cache_path)else:image = Image.open(img_path).convert('RGB')if transform is not None:image = transform(image)torch.save(image, cache_path)  # 保存到缓存目录self.imgs.append(image)self.labels.append(label)

优点：不占用内存，适合大数据集；
缺点：首次加载需要时间生成缓存，且占用磁盘空间（与原始图像相当）。

4.3 优化 3：处理不均衡数据集（避免模型偏向多数类）

食品数据集中常存在类别不均衡问题（如 “汉堡” 样本有 1000 张，“鱼子酱” 样本只有 100 张），若不处理，模型会偏向预测多数类（如预测所有样本为 “汉堡”），导致评估指标（如 F1-score）偏低。

解决方法是使用WeightedRandomSampler，根据类别权重随机采样，让每个类别被选中的概率相同。

4.3.1 实现步骤

1. 计算每个类别的样本数量；
2. 计算每个类别的权重（权重 = 总样本数 / 类别样本数）；
3. 为每个样本分配对应的类别权重；
4. 将WeightedRandomSampler传入DataLoader的sampler参数。

4.3.2 代码实现

python

from torch.utils.data import WeightedRandomSampler# 1. 计算训练集每个类别的样本数量
train_labels = train_dataset.labels  # 训练集所有标签
classes = list(set(train_labels))  # 所有类别
class_count = {cls: train_labels.count(cls) for cls in classes}  # 类别-样本数字典
print("类别样本数：", class_count)  # 输出：{0:1000, 1:100}（假设0=汉堡，1=鱼子酱）# 2. 计算每个类别的权重（总样本数/类别样本数）
total_samples = len(train_labels)
class_weight = {cls: total_samples / count for cls, count in class_count.items()}
print("类别权重：", class_weight)  # 输出：{0: 0.1, 1: 1.0}（总样本数1100）# 3. 为每个样本分配权重（样本权重=其类别权重）
sample_weights = [class_weight[label] for label in train_labels]# 4. 创建WeightedRandomSampler
sampler = WeightedRandomSampler(weights=sample_weights,  # 样本权重列表num_samples=len(sample_weights),  # 采样数量（等于总样本数）replacement=True  # 是否允许重复采样（True表示允许，确保每个epoch采样数量足够）
)# 5. 将sampler传入DataLoader（注意：此时shuffle必须设为False，否则sampler无效）
train_loader_balanced = DataLoader(train_dataset,batch_size=32,shuffle=False,  # 必须设为Falsesampler=sampler,  # 传入加权采样器num_workers=4,pin_memory=True
)

效果：每个 epoch 中，“汉堡” 和 “鱼子酱” 样本被选中的概率相同，模型会更关注少数类，提升少数类的预测准确率。

4.4 优化 4：使用混合精度加载（节省内存）

对于 16 位 GPU（如 RTX 3090、A100），可以使用混合精度加载数据，将图像张量从float32转为float16，减少内存占用。

代码实现

python

# 在预处理流水线中添加类型转换
data_transforms['train'] = transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(p=0.5),transforms.ToTensor(),transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),transforms.ConvertImageDtype(torch.float16)  # 转为float16
])# 验证集同理
data_transforms['valid'].transforms.append(transforms.ConvertImageDtype(torch.float16))

优点：内存占用减少一半（float32每个数值占 4 字节，float16占 2 字节），可支持更大的batch_size；
缺点：部分模型可能对float16的精度敏感，导致性能轻微下降（可通过混合精度训练弥补）。

4.5 优化 5：数据可视化（验证预处理是否正确）

预处理后的图像是否正确？标签是否与图像匹配？这些问题可以通过数据可视化来验证，避免因预处理错误导致模型训练无效。

代码实现（使用 matplotlib）

python

import matplotlib.pyplot as plt# 定义类别名称（假设0=汉堡，1=披萨，2=牛排）
class_names = ["汉堡", "披萨", "牛排"]# 从训练加载器中获取一个批次
for images, labels in train_loader:# 转换为CPU张量（若在GPU上）images = images.cpu()labels = labels.cpu()# 绘制4×4的图像网格fig, axes = plt.subplots(4, 4, figsize=(12, 12))axes = axes.flatten()for img, label, ax in zip(images[:16], labels[:16], axes):# 转换张量形状：C×H×W → H×W×Cimg = img.permute(1, 2, 0)# 反归一化（将像素值从[-1,1]转回[0,1]，方便显示）mean = torch.tensor([0.485, 0.456, 0.406])std = torch.tensor([0.229, 0.224, 0.225])img = img * std + mean# 限制像素值范围（避免因反归一化导致的数值超出[0,1]）img = torch.clamp(img, 0, 1)# 显示图像和标签ax.imshow(img.numpy())ax.set_title(f"类别：{class_names[label.item()]}")ax.axis('off')  # 隐藏坐标轴plt.tight_layout()plt.show()break  # 仅显示一个批次

作用：通过可视化，我们可以直观地看到：

• 预处理后的图像是否变形（如 Resize 是否正确）；
• 数据增强是否生效（如随机翻转、裁剪是否正确）；
• 标签是否与图像匹配（如标签为 “汉堡” 的图像是否确实是汉堡）。

五、常见问题与解决方案（新手避坑指南）

在使用本文代码时，新手可能会遇到各种问题。本节整理了 10 个最常见的问题，并提供详细的解决方案，帮助你快速排查错误。

问题 1：NameError: name 'Dataset' is not defined

原因：未从torch.utils.data导入Dataset类。
解决方案：确保导入语句正确：

python

from torch.utils.data import Dataset, DataLoader

问题 2：FileNotFoundError: [Errno 2] No such file or directory: 'train_labels.txt'

原因：标签文件路径错误（如文件不存在、路径拼写错误）。
解决方案：

1. 检查file_path是否正确，使用绝对路径（如"D:/food_data/train_labels.txt"）；
2. 使用os.path.abspath(file_path)查看文件的绝对路径，确认是否存在；
3. 若使用相对路径，确保标签文件在当前工作目录下（可通过os.getcwd()查看当前工作目录）。

问题 3：RuntimeError: stack expects each tensor to be equal size, but got [3,224,224] and [1,224,224]

原因：批次中存在通道数不一致的图像（如 RGB 三通道和灰度单通道）。
解决方案：在__getitem__中强制将图像转为 RGB：

python

image = Image.open(img_path).convert('RGB')

问题 4：RuntimeError: expected scalar type Long but found Int

原因：标签张量类型不是int64（长整型），而是int32（整型）。
解决方案：创建标签张量时指定dtype=torch.int64：

python

label_tensor = torch.tensor(label, dtype=torch.int64)

问题 5：BrokenPipeError: [Errno 32] Broken pipe（Windows 系统）

原因：Windows 下多进程加载（num_workers>0）兼容性问题。
解决方案：

1. 将num_workers设为 0（仅主进程加载，速度慢但稳定）；
2. 在脚本开头添加if __name__ == "__main__":，将所有代码放入该块中；
3. 更新 PyTorch 到 1.8 + 版本（修复了部分 Windows 多进程 bug）。

问题 6：OOMError: CUDA out of memory

原因：GPU 内存不足，通常是batch_size过大导致。
解决方案：

1. 减小batch_size（如从 32 改为 16、8）；
2. 使用混合精度加载（float16），减少内存占用；
3. 关闭pin_memory（仅在内存不足时使用，会降低速度）；
4. 使用更小的图像尺寸（如从 224×224 改为 112×112）。

问题 7：模型训练时 loss 不下降

原因：预处理错误（如未归一化、归一化参数错误）。
解决方案：

1. 检查是否添加了Normalize操作；
2. 若使用预训练模型，确保Normalize的均值方差与预训练数据一致（如 ImageNet 的参数）；
3. 验证预处理后的图像像素值是否在合理范围（如归一化后应为 [-1,1] 左右）。

问题 8：数据加载速度过慢（每个 epoch 耗时过长）

原因：num_workers设置过小、未使用pin_memory、未缓存数据。
解决方案：

1. 将num_workers设为 CPU 核心数的 1~2 倍；
2. 使用 GPU 时开启pin_memory=True；
3. 对大数据集使用磁盘缓存（如.pt文件）；
4. 确保图像文件存储在 SSD 上（SSD 读取速度远快于 HDD）。

问题 9：标签与图像不匹配

原因：标签文件格式错误（如行分割符不是空格、标签顺序错误）。
解决方案：

1. 打开标签文件，检查每行格式是否为 “图像路径 标签”（如 “food/001.jpg 0”）；
2. 若使用其他分割符（如逗号），修改split(' ')为split(',')；
3. 通过数据可视化验证标签与图像是否匹配。

问题 10：读取图像时出现 “OSError: cannot identify image file”

原因：图像文件损坏或格式不支持。
解决方案：

1. 验证图像路径是否正确，尝试用 PIL 手动打开该图像（Image.open("path/to/img.jpg")）；
2. 删除损坏的图像文件，或用工具修复；
3. 在__getitem__中添加异常处理，跳过损坏的图像。

六、总结：数据加载的核心原则与学习建议

通过本文的学习，我们从基础原理到实战优化，完整掌握了 PyTorch 数据加载的全流程。最后，总结几个核心原则和学习建议，帮助你在后续项目中灵活应用。