大家好，我是南木。数据加载是深度学习训练的“第一环”，也是最容易被忽视的性能瓶颈。很多人花几周调参优化模型，却因为数据加载“喂不饱”GPU，导致整体训练效率折损50%以上。更关键的是，自定义数据集的代码质量直接决定了项目的可维护性——不少团队因为前期数据加载逻辑混乱，后期迭代时不得不推倒重写。

这篇文章将从“基础原理→自定义实现→性能优化→高级技巧”全流程拆解PyTorch数据加载。包含图像/文本/多模态数据集的自定义模板、多线程参数调优指南、10个高频BUG修复方案，附带可直接复用的代码模板和性能测试工具。

同时这里还给大家整理了一份适合零基础入门学习的深度学习资料包 需要的同学扫码自取

一、先搞懂：PyTorch数据加载的核心逻辑（2个核心组件）

PyTorch的数据加载体系主要依赖Dataset和DataLoader两个组件，二者分工明确，就像“仓库管理员”和“配送员”的关系——Dataset负责数据的“存取管理”，DataLoader负责数据的“高效配送”。

1. Dataset：数据的“仓库管理员”

Dataset的核心作用是定义数据的读取逻辑，告诉PyTorch“数据在哪、怎么读、怎么返回”。它本质是一个抽象类，我们需要继承它并实现两个核心方法：

• __len__()：返回数据集的总样本数（告诉仓库有多少货）；
• __getitem__(idx)：根据索引idx返回单条样本（根据订单取货）。

（1）最简化示例：自定义一个数值数据集

import torch
from torch.utils.data import Dataset

# 自定义数据集：生成y = 2x + 3的样本
class SimpleDataset(Dataset):
    def __init__(self, num_samples=100):
        # 初始化：生成数据（相当于“整理仓库货物”）
        self.x = torch.randn(num_samples, 1)  # 输入特征
        self.y = 2 * self.x + 3 + torch.randn(num_samples, 1) * 0.1  # 带噪声的标签
    
    def __len__(self):
        # 返回数据集大小
        return len(self.x)
    
    def __getitem__(self, idx):
        # 根据索引返回单条样本
        return self.x[idx], self.y[idx]

# 测试数据集
dataset = SimpleDataset(num_samples=10)
print(f"数据集大小：{len(dataset)}")  # 输出10
x, y = dataset[0]  # 取第0个样本
print(f"第0个样本：x={x.item():.4f}, y={y.item():.4f}")

（2）核心原则：__getitem__必须返回“可张量化”的数据

__getitem__的返回值通常是(特征, 标签)的元组，要求：

• 特征：可以是numpy.ndarray、PIL.Image或torch.Tensor（最终会被DataLoader转为Tensor）；
• 标签：通常是整数（分类）或浮点数（回归）；
• 禁止返回“可变长度”的列表（如不同长度的文本），需先做padding处理。

2. DataLoader：数据的“高效配送员”

Dataset只解决了“怎么取单条数据”，而DataLoader则解决了“如何高效批量配送”——它会自动完成批量拼接（batch）、打乱（shuffle）、多线程加载（num_workers） 等操作，核心目标是“让GPU在训练时永远不缺数据”。

（1）基础用法：从Dataset到DataLoader

from torch.utils.data import DataLoader

# 1. 实例化数据集
dataset = SimpleDataset(num_samples=1000)

# 2. 实例化DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,  # 每批32个样本
    shuffle=True,   # 训练时打乱数据
    num_workers=4,  # 4个线程并行加载
    pin_memory=True # 锁存内存，加速GPU数据传输
)

# 3. 迭代使用（训练时的典型用法）
for epoch in range(5):
    for batch_x, batch_y in dataloader:
        # batch_x形状：[32, 1]，batch_y形状：[32, 1]
        # 模型训练逻辑...
        print(f"Batch x shape: {batch_x.shape}, Batch y shape: {batch_y.shape}")
        break  # 只打印第一个batch

（2）核心参数解析（性能优化的关键）

DataLoader的参数直接影响加载速度，新手最容易在num_workers和pin_memory上踩坑，这里逐一拆解：

参数名	作用	新手建议
batch_size	每批样本数	训练时设32/64（根据GPU显存调整），验证时设128/256
shuffle	是否打乱数据	训练时True，验证时False（保持结果可复现）
num_workers	并行加载的线程数	Linux设为CPU核心数的1-2倍，Windows设为0（避免BrokenPipeError）
pin_memory	是否锁存内存	当使用GPU时设为True（加速数据从CPU到GPU的传输）
drop_last	最后一批样本不足batch_size时是否丢弃	训练时True（避免batch_size不一致导致的BatchNorm问题），验证时False
prefetch_factor	每个线程预加载的批次数	PyTorch 1.7+支持，设为2（提前加载下一批，避免GPU等待）

3. 性能瓶颈的根源：为什么数据加载会“卡壳”？

很多人发现“GPU利用率上不去”，本质是数据加载速度跟不上GPU的计算速度，导致GPU经常“闲等”数据。常见原因有三个：

1. 单线程加载太慢：Python的GIL锁导致单线程读取数据效率低，尤其当数据需要解压（如JPEG）、预处理（如Resize）时；
2. 数据预处理耗时：将数据增强（如RandomCrop、Flip）放在__getitem__中，若操作复杂，单条样本处理时间过长；
3. IO瓶颈：数据存放在机械硬盘（HDD）上，或网络存储（NAS）延迟高，读取速度慢于GPU计算速度。

解决思路：用多线程（num_workers）并行加载，将预处理“前移”（如提前生成LMDB文件），或用SSD/内存加速IO。

二、实战：3类自定义Dataset模板（覆盖90%场景）

自定义Dataset是数据加载的核心，不同数据类型（图像、文本、多模态）的实现逻辑不同。以下是工业项目中最常用的3类模板，可直接复用。

1. 图像数据集（最常用，以分类任务为例）

图像数据集通常的目录结构是“按类别分文件夹”，如：

data/
├── cat/
│   ├── cat1.jpg
│   ├── cat2.jpg
│   └── ...
└── dog/
    ├── dog1.jpg
    ├── dog2.jpg
    └── ...

（1）自定义图像数据集代码

import os
import cv2
import torch
from torch.utils.data import Dataset
from torchvision import transforms

class ImageClassificationDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        """
        图像分类数据集
        :param data_dir: 数据根目录（含类别子文件夹）
        :param transform: 数据增强/预处理 transforms
        """
        self.data_dir = data_dir
        self.transform = transform
        
        # 1. 遍历目录，获取所有图像路径和标签
        self.image_paths = []
        self.labels = []
        self.classes = os.listdir(data_dir)  # 获取类别名（如["cat", "dog"]）
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}  # 类别到索引的映射
        
        for cls in self.classes:
            cls_dir = os.path.join(data_dir, cls)
            for img_name in os.listdir(cls_dir):
                # 过滤非图像文件
                if img_name.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
                    img_path = os.path.join(cls_dir, img_name)
                    self.image_paths.append(img_path)
                    self.labels.append(self.class_to_idx[cls])
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        # 1. 读取图像（用cv2或PIL，这里用cv2，返回BGR格式）
        img_path = self.image_paths[idx]
        img = cv2.imread(img_path)
        if img is None:
            # 处理损坏图像（返回前一个样本，避免报错）
            return self.__getitem__((idx + 1) % len(self))
        # 转为RGB格式（cv2默认BGR，与PyTorch的预处理一致）
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 2. 数据增强/预处理
        if self.transform is not None:
            img = self.transform(img)
        
        # 3. 返回图像和标签
        label = self.labels[idx]
        return img, label, img_path  # 额外返回路径，方便调试

（2）使用示例（含数据增强）

# 定义数据增强（训练时用，验证时不用Random操作）
train_transform = transforms.Compose([
    transforms.ToPILImage(),  # cv2读取的是ndarray，需转为PIL才能用torchvision.transforms
    transforms.Resize((256, 256)),
    transforms.RandomCrop((224, 224)),  # 随机裁剪
    transforms.RandomHorizontalFlip(p=0.5),  # 随机水平翻转
    transforms.ToTensor(),  # 转为Tensor（HWC→CHW，值归一化到[0,1]）
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化（ImageNet均值）
])

# 验证时的预处理（无随机操作）
val_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 实例化数据集
train_dataset = ImageClassificationDataset(
    data_dir="data/train",
    transform=train_transform
)
val_dataset = ImageClassificationDataset(
    data_dir="data/val",
    transform=val_transform
)

# 实例化DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=True
)
val_loader = DataLoader(
    val_dataset,
    batch_size=64,
    shuffle=False,
    num_workers=2,
    pin_memory=True
)

# 测试
for img, label, path in train_loader:
    print(f"Img shape: {img.shape}, Label shape: {label.shape}, Path: {path[0]}")
    # Img shape: torch.Size([32, 3, 224, 224]), Label shape: torch.Size([32])
    break

（3）避坑点：

• 图像格式转换：cv2默认BGR，PIL默认RGB，必须统一为RGB（与预处理的均值/std匹配）；
• 损坏图像处理：cv2.imread()读取损坏图像会返回None，需加异常处理；
• 数据增强顺序：先Resize到大于目标尺寸（如256→224），再RandomCrop，避免图像拉伸变形。

2. 文本数据集（以文本分类为例）

文本数据集通常是“CSV文件”或“文本文件”，每行包含“文本+标签”，如：

text,label
"这只猫很可爱",0
"今天天气很好",1
"这部电影真难看",2

（1）自定义文本数据集代码

import pandas as pd
import torch
from torch.utils.data import Dataset
from transformers import BertTokenizer  # 用Hugging Face的Tokenizer做文本预处理

class TextClassificationDataset(Dataset):
    def __init__(self, csv_path, tokenizer, max_len=512):
        """
        文本分类数据集
        :param csv_path: CSV文件路径（含text和label列）
        :param tokenizer: 分词器（如BertTokenizer）
        :param max_len: 文本最大长度（超过截断，不足padding）
        """
        self.data = pd.read_csv(csv_path)
        self.texts = self.data["text"].tolist()
        self.labels = self.data["label"].tolist()
        self.tokenizer = tokenizer
        self.max_len = max_len
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        # 1. 获取文本和标签
        text = str(self.texts[idx])  # 确保是字符串，避免NaN
        label = self.labels[idx]
        
        # 2. 文本分词（BERT风格，含padding和truncation）
        encoding = self.tokenizer(
            text,
            add_special_tokens=True,  # 加[CLS]和[SEP]
            max_length=self.max_len,
            padding="max_length",  # 不足max_len则padding
            truncation=True,  # 超过max_len则截断
            return_attention_mask=True,  # 返回attention mask
            return_tensors="pt"  # 返回Tensor
        )
        
        # 3. 整理输出（去除batch维度，因为DataLoader会批量拼接）
        input_ids = encoding["input_ids"].flatten()
        attention_mask = encoding["attention_mask"].flatten()
        
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "label": torch.tensor(label, dtype=torch.long)
        }

（2）使用示例（基于BERT分词器）

# 加载分词器（需先安装transformers：pip install transformers）
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")

# 实例化数据集
train_dataset = TextClassificationDataset(
    csv_path="data/train.csv",
    tokenizer=tokenizer,
    max_len=128
)
val_dataset = TextClassificationDataset(
    csv_path="data/val.csv",
    tokenizer=tokenizer,
    max_len=128
)

# 实例化DataLoader
train_loader = DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=2,
    pin_memory=True
)

# 测试
for batch in train_loader:
    print(f"Input_ids shape: {batch['input_ids'].shape}")  # [16, 128]
    print(f"Attention_mask shape: {batch['attention_mask'].shape}")  # [16, 128]
    print(f"Label shape: {batch['label'].shape}")  # [16]
    break

（3）避坑点：

• 文本格式处理：需处理NaN、空字符串等异常，避免分词器报错；
• max_len设置：根据模型最大输入长度（如BERT-base最大512）和文本实际长度设置，过长会浪费内存；
• Tokenizer复用：Tokenizer应在Dataset外实例化，避免重复加载（节省内存）。

3. 多模态数据集（图像+文本，以图文匹配为例）

多模态数据集需要同时加载不同类型的数据（如商品图像+商品描述），核心是“保持样本对齐”（同一索引对应同一样本的不同模态）。

（1）自定义多模态数据集代码

import os
import cv2
import pandas as pd
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from transformers import BertTokenizer

class ImageTextDataset(Dataset):
    def __init__(self, csv_path, img_dir, tokenizer, img_transform=None, max_len=512):
        """
        图文多模态数据集
        :param csv_path: CSV路径（含img_name、text、label列）
        :param img_dir: 图像根目录
        :param tokenizer: 文本分词器
        :param img_transform: 图像预处理
        :param max_len: 文本最大长度
        """
        self.data = pd.read_csv(csv_path)
        self.img_dir = img_dir
        self.tokenizer = tokenizer
        self.img_transform = img_transform
        self.max_len = max_len
        
        # 检查必要的列是否存在
        required_cols = ["img_name", "text", "label"]
        for col in required_cols:
            if col not in self.data.columns:
                raise ValueError(f"CSV文件缺少必要列：{col}")
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        # 1. 读取图像
        img_name = self.data.iloc[idx]["img_name"]
        img_path = os.path.join(self.img_dir, img_name)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        if self.img_transform is not None:
            img = self.img_transform(img)
        
        # 2. 读取并处理文本
        text = str(self.data.iloc[idx]["text"])
        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            padding="max_length",
            truncation=True,
            return_attention_mask=True,
            return_tensors="pt"
        )
        input_ids = encoding["input_ids"].flatten()
        attention_mask = encoding["attention_mask"].flatten()
        
        # 3. 读取标签
        label = self.data.iloc[idx]["label"]
        
        return {
            "image": img,
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "label": torch.tensor(label, dtype=torch.long)
        }

（2）使用示例

# 初始化图像预处理和文本分词器
img_transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")

# 实例化数据集
dataset = ImageTextDataset(
    csv_path="data/multimodal_train.csv",
    img_dir="data/images",
    tokenizer=tokenizer,
    img_transform=img_transform,
    max_len=128
)

# 实例化DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=8,
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

# 测试
for batch in dataloader:
    print(f"Image shape: {batch['image'].shape}")  # [8, 3, 224, 224]
    print(f"Input_ids shape: {batch['input_ids'].shape}")  # [8, 128]
    print(f"Label shape: {batch['label'].shape}")  # [8]
    break

（3）避坑点：

• 样本对齐：确保CSV中的img_name与图像文件名完全一致（包括大小写、后缀）；
• 数据类型统一：不同模态的数据预处理要匹配模型输入（如图像转为CHW格式，文本转为input_ids）；
• 加载顺序：多模态数据加载耗时更长，建议num_workers设为CPU核心数的1倍，避免内存溢出。

三、性能优化：从“卡壳”到“喂饱GPU”的5个实战技巧

数据加载的终极目标是“让GPU利用率稳定在90%以上”。以下5个技巧是我在工业项目中验证过的“性能倍增器”，从IO、多线程、预处理三个维度全面优化。

1. 技巧1：用SSD/内存盘加速IO（最直接的硬件优化）

数据加载的第一瓶颈往往是磁盘IO速度——机械硬盘（HDD）的读取速度通常只有100-200MB/s，而SSD可达500-3000MB/s，内存盘（RAM Disk）更是能到10GB/s以上。

（1）实操方案：

• 短期：将数据集复制到SSD上，训练时从SSD读取；
• 长期：对于超大规模数据集（如百万级图像），使用分布式存储（如Ceph）或GPU直接访问（如NVIDIA Magnum IO）；
• 应急：在Linux上用tmpfs创建内存盘，将小数据集加载到内存中：

  # 创建10GB内存盘
  sudo mkdir /mnt/ramdisk
  sudo mount -t tmpfs -o size=10G tmpfs /mnt/ramdisk
  # 将数据集复制到内存盘
  cp -r data/train /mnt/ramdisk/
  # 训练时从内存盘读取
  

（2）效果对比：

存储类型	读取速度	1000张224x224图像加载时间	GPU利用率
机械硬盘（HDD）	150MB/s	8.2秒	30-50%
SATA SSD	500MB/s	2.5秒	70-80%
NVMe SSD	3000MB/s	0.4秒	85-95%
内存盘	10GB/s	0.1秒	90-98%

2. 技巧2：num_workers参数调优（多线程加载的黄金法则）

num_workers是控制并行加载的核心参数，但并非“越大越好”——过多的线程会导致CPU上下文切换频繁，反而降低速度。

（1）调优原则：

• Linux系统：num_workers = CPU核心数 × 1 ~ 2（如8核CPU设8-16）；
• Windows系统：num_workers = 0（Windows的多线程支持差，容易报BrokenPipeError）；
• 验证方法：用timeit测试不同num_workers的加载时间，选择最优值：

  import timeit
  
  def test_dataloader(num_workers):
      dataloader = DataLoader(
          train_dataset,
          batch_size=32,
          shuffle=True,
          num_workers=num_workers,
          pin_memory=True
      )
      # 迭代一次，测试加载时间
      for batch in dataloader:
          pass
  
  # 测试不同num_workers
  for workers in [0, 2, 4, 8, 16]:
      time = timeit.timeit(lambda: test_dataloader(workers), number=3)
      print(f"num_workers={workers}, 平均时间：{time/3:.2f}秒")
  

（2）常见误区：

• ❌ 错误：“num_workers设成32，速度肯定最快”——8核CPU设32会导致线程争抢，速度反而比8慢30%；
• ✅ 正确：根据CPU核心数调整，优先用htop查看CPU利用率，确保加载时CPU利用率在70-80%。

3. 技巧3：数据预处理“前移”（避免在__getitem__中做 heavy 操作）

很多人将“Resize、Normalize”等耗时操作放在__getitem__中，导致单条样本处理时间过长。正确的做法是将预处理“前移”到数据准备阶段，提前生成处理好的数据。

（1）实操方案：

1. 提前Resize：用脚本批量将图像Resize到目标尺寸，保存为新文件；
2. 生成LMDB文件：将图像和标签存入LMDB数据库（键值对存储，读取速度比文件系统快3-5倍）；
3. 使用DALI加速：NVIDIA DALI库可将预处理卸载到GPU，比CPU预处理快10倍以上。

（2）LMDB数据集示例：

import lmdb
import cv2
import pickle
import torch
from torch.utils.data import Dataset

class LMDBImageDataset(Dataset):
    def __init__(self, lmdb_path, transform=None):
        self.lmdb_path = lmdb_path
        self.transform = transform
        # 打开LMDB数据库
        self.env = lmdb.open(
            lmdb_path,
            max_readers=1,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False
        )
        # 读取元数据
        with self.env.begin(write=False) as txn:
            self.length = pickle.loads(txn.get(b"__length__"))
            self.keys = pickle.loads(txn.get(b"__keys__"))
    
    def __getitem__(self, idx):
        key = self.keys[idx]
        with self.env.begin(write=False) as txn:
            value = txn.get(key.encode())
            # 解析数据（假设存储的是(img, label)的pickle对象）
            img, label = pickle.loads(value)
        # 预处理
        if self.transform is not None:
            img = self.transform(img)
        return img, label
    
    def __len__(self):
        return self.length

（3）效果：预处理前移后，__getitem__的处理时间从0.02秒/样本降至0.005秒/样本，加载速度提升4倍。

4. 技巧4：用Albumentations替代torchvision.transforms（数据增强加速）

数据增强是提升精度的关键，但torchvision.transforms的速度较慢，尤其在多线程加载时。Albumentations是一个专门的图像增强库，速度比torchvision快2-3倍，且支持更多增强操作。

（1）使用示例：

# 安装Albumentations：pip install albumentations
import albumentations as A
from albumentations.pytorch import ToTensorV2

# 定义Albumentations增强流水线（比torchvision快3倍）
train_transform = A.Compose([
    A.Resize(height=256, width=256),
    A.RandomCrop(height=224, width=224),
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2()  # 直接转为PyTorch Tensor
])

# 自定义Dataset中使用Albumentations（无需ToPILImage，直接处理ndarray）
class ImageDatasetWithAlbumentations(Dataset):
    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # 直接用Albumentations处理ndarray
        augmented = train_transform(image=img)
        img = augmented["image"]
        label = self.labels[idx]
        return img, label

（2）核心优势：

• 速度快：基于OpenCV实现，处理ndarray比torchvision的PIL操作快；
• 支持更多操作：如CutMix、Mosaic、GridDistortion等高级增强；
• 多模态支持：可同时处理图像、掩码（如语义分割）、关键点（如目标检测）。

5. 技巧5：监控工具实时调优（找到瓶颈的“放大镜”）

优化数据加载的前提是“找到瓶颈”，以下工具可实时监控加载速度和资源占用：

（1）GPU利用率监控：nvidia-smi

# 实时监控GPU利用率（每秒刷新一次）
watch -n 1 nvidia-smi

• 若GPU利用率低于70%，且gpu_utils列显示“空闲”，说明数据加载慢；
• 若Memory-Usage接近满，说明batch_size过大，需减小。

（2）CPU和IO监控：htop + iotop

# 监控CPU和内存占用
htop
# 监控磁盘IO占用
sudo iotop

• 若htop显示CPU利用率接近100%，说明num_workers过多；
• 若iotop显示磁盘IO利用率接近100%，说明磁盘速度是瓶颈，需换SSD。

（3）PyTorch内置监控：DataLoader迭代时间

# 监控每个batch的加载时间
import time

start_time = time.time()
for batch_idx, (img, label) in enumerate(train_loader):
    # 计算加载时间
    batch_time = time.time() - start_time
    print(f"Batch {batch_idx}, 加载时间：{batch_time:.4f}秒")
    # 模拟模型训练（占用GPU）
    with torch.no_grad():
        output = model(img.cuda())
    # 重置计时器
    start_time = time.time()

• 若加载时间 > 模型计算时间，说明数据加载是瓶颈；
• 理想状态：加载时间 < 模型计算时间（GPU永远不等待）。

四、10个高频BUG修复：从报错到解决（新手必看）

数据加载的BUG看似五花八门，实则多是“路径问题”“格式问题”“参数问题”三类。以下是10个新手最常踩的坑，每个都附“错误示例+原因分析+解决方案”。

1. BUG1：FileNotFoundError（最常见的路径错误）

错误示例：

# 数据集目录结构：data/train/cat/cat1.jpg
dataset = ImageClassificationDataset(data_dir="data/train/cat")  # 错误：传入了类别子目录

报错：FileNotFoundError: [Errno 2] No such file or directory: 'data/train/cat/cat'

原因分析：data_dir应传入“根目录”（含类别子文件夹），而非直接传入类别子目录。

解决方案：传入正确的根目录：

dataset = ImageClassificationDataset(data_dir="data/train")  # 正确：根目录含cat和dog子文件夹

2. BUG2：BrokenPipeError（Windows多线程加载错误）

错误示例：

# Windows系统下设置num_workers=4
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)  # 错误

报错：BrokenPipeError: [Errno 32] Broken pipe

原因分析：Windows的多进程支持不完善，num_workers>0时容易出现管道断裂。

解决方案：Windows系统下num_workers=0：

dataloader = DataLoader(dataset, batch_size=32, num_workers=0)  # 正确

3. BUG3：RuntimeError（batch_size不一致导致的BatchNorm错误）

错误示例：

# 数据集大小1001，batch_size=32，最后一批1个样本
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False)  # 错误

报错：RuntimeError: running_mean should contain 32 elements not 1

原因分析：最后一批样本不足batch_size，导致BatchNorm层的统计量计算错误。

解决方案：训练时drop_last=True，丢弃最后一批：

dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=True)  # 正确

4. BUG4：TypeError（数据类型不匹配）

错误示例：

# __getitem__返回列表而非Tensor
def __getitem__(self, idx):
    x = [1, 2, 3]  # 错误：返回列表
    y = 0
    return x, y

报错：TypeError: can't convert list to Tensor

原因分析：DataLoader无法将列表直接转为Tensor，需返回ndarray或Tensor。

解决方案：返回ndarray或Tensor：

def __getitem__(self, idx):
    x = np.array([1, 2, 3])  # 正确：返回ndarray
    y = 0
    return x, y

5. BUG5：内存泄漏（长期训练内存越来越高）

错误示例：

# 在__getitem__中创建大量临时变量
def __getitem__(self, idx):
    img_path = self.image_paths[idx]
    img = cv2.imread(img_path)
    # 创建大量临时数组
    temp1 = np.zeros((224, 224))
    temp2 = np.zeros((224, 224))
    # ... 更多临时变量
    return img, label

现象：训练几小时后，内存占用从8GB升至32GB，最终OOM。

原因分析：多线程加载时，临时变量未及时释放，导致内存泄漏。

解决方案：

1. 减少临时变量，尽量复用数组；
2. 在Linux上用valgrind检测内存泄漏；
3. 升级PyTorch到1.10+（修复了多个DataLoader内存泄漏问题）。

6. BUG6：数据加载顺序错乱（shuffle=True但结果不随机）

错误示例：

# 多次实例化Dataset，每次都会重新生成图像路径（导致shuffle无效）
for epoch in range(10):
    dataset = ImageClassificationDataset(data_dir="data/train")  # 错误：每次都重新创建
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

现象：每次epoch的加载顺序相同，shuffle未生效。

原因分析：每次重新创建Dataset时，image_paths的顺序相同，shuffle只是在相同顺序上打乱，导致随机性不足。

解决方案：只实例化一次Dataset，重复使用：

# 正确：只实例化一次Dataset
dataset = ImageClassificationDataset(data_dir="data/train")
for epoch in range(10):
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True)  # 重复使用Dataset

7. BUG7：cv2.imread返回None（图像损坏或路径错误）

错误示例：

def __getitem__(self, idx):
    img_path = self.image_paths[idx]
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 错误：若img为None，会报错

报错：error: (-215:Assertion failed) !_src.empty() in function 'cvtColor'

原因分析：图像文件损坏或路径错误，导致cv2.imread返回None。

解决方案：添加异常处理，跳过损坏图像：

def __getitem__(self, idx):
    img_path = self.image_paths[idx]
    img = cv2.imread(img_path)
    if img is None:
        # 返回下一个样本，避免报错
        return self.__getitem__((idx + 1) % len(self))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img, label

8. BUG8：pin_memory=True导致内存不足

错误示例：

# 小内存机器上启用pin_memory=True
dataloader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)  # 错误

报错：RuntimeError: CUDA out of memory

原因分析：pin_memory=True会将数据锁存在内存中，占用更多内存，小内存机器容易OOM。

解决方案：

1. 小内存机器（<16GB）禁用pin_memory；
2. 减小batch_size，释放内存；
3. 使用torch.cuda.empty_cache()定期清理GPU缓存。

9. BUG9：多进程加载时数据增强不一致

错误示例：

# 在Dataset中实例化随机增强对象（多进程时会产生相同随机种子）
class ImageDataset(Dataset):
    def __init__(self):
        self.flip = transforms.RandomHorizontalFlip(p=0.5)  # 错误：多进程共享同一随机对象
    
    def __getitem__(self, idx):
        img = cv2.imread(img_path)
        img = self.flip(transforms.ToPILImage()(img))
        return img

现象：多进程加载时，不同进程的增强结果相同，随机性不足。

原因分析：多进程会复制父进程的随机状态，导致不同进程的增强结果一致。

解决方案：在__getitem__中创建随机增强对象，或设置不同随机种子：

def __getitem__(self, idx):
    # 正确：在__getitem__中创建增强对象
    flip = transforms.RandomHorizontalFlip(p=0.5)
    img = cv2.imread(img_path)
    img = flip(transforms.ToPILImage()(img))
    return img

10. BUG10：分布式训练时数据重复加载

错误示例：

# 分布式训练时未用DistributedSampler，导致每个进程加载全部数据
train_sampler = torch.utils.data.RandomSampler(train_dataset)  # 错误
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    sampler=train_sampler,
    num_workers=4
)

现象：4个GPU进程各加载全部数据，导致训练样本重复4次，精度不收敛。

原因分析：分布式训练时，需用DistributedSampler将数据分片到不同进程，避免重复。

解决方案：使用DistributedSampler：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)  # 正确
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    sampler=train_sampler,
    num_workers=4
)
# 训练时设置epoch
for epoch in range(10):
    train_sampler.set_epoch(epoch)  # 确保每个epoch的分片不同
    for batch in train_loader:
        # 训练逻辑...

五、学习路径：从入门到精通数据加载（3个阶段）

数据加载看似简单，实则需要“工程能力+性能优化意识”的结合。以下是针对不同阶段的学习路径，帮你系统提升。

1. 入门阶段（1-2周）：掌握基础实现

目标：能自定义数据集，正确使用DataLoader加载数据。
核心任务：

1. 理解Dataset和DataLoader的分工，实现图像/文本数据集；
2. 掌握常用数据预处理（Resize、Normalize、分词）；
3. 解决“路径错误”“格式不匹配”等基础BUG。
   推荐资源：

• PyTorch官方教程：Data Loading and Processing Tutorial；
• 实战项目：实现MNIST手写数字分类的数据加载（含数据增强）。

2. 进阶阶段（1-2个月）：性能优化与复杂场景

目标：优化加载速度，掌握多模态、分布式等复杂场景。
核心任务：

1. 调优num_workers、pin_memory等参数，提升GPU利用率；
2. 实现多模态数据集（图像+文本+音频）；
3. 掌握LMDB、DALI等加速工具的使用；
4. 实现分布式训练的数据加载（DistributedSampler）。
   推荐资源：

• Albumentations文档：Albumentations Documentation；
• NVIDIA DALI教程：DALI Getting Started。

3. 专家阶段（2-3个月）：底层优化与定制化

目标：深入DataLoader底层，定制化加载逻辑。
核心任务：

1. 阅读PyTorch DataLoader源码，理解多进程加载原理；
2. 实现定制化Sampler（如按样本难度采样、类别平衡采样）；
3. 优化超大规模数据集（1000万+样本）的加载策略；
4. 结合硬件特性（如NVMe SSD、GPU Direct Storage）优化IO。
   推荐资源：

• PyTorch源码：torch/utils/data；
• 论文：Optimizing Data Loading for Deep Learning with PyTorch。

我是南木，专注AI技术实战与学习规划。后续会分享更多PyTorch核心知识点（如模型构建、分布式训练），关注我，一起少走弯路，高效进阶深度学习！