数据加载与预处理

高效的数据加载与预处理是深度学习训练流水线中至关重要的一环。PyTorch 提供了 torch.utils.data 模块，包含 Dataset 和 DataLoader 两个核心抽象，配合 torchvision.transforms 可以构建灵活、高效的数据处理管道。

---

Dataset 抽象类

PyTorch 中所有数据集都应继承 torch.utils.data.Dataset 并实现两个核心方法：

方法	说明	必须实现
__getitem__(self, index)	根据索引获取单个样本	是
__len__(self)	返回数据集大小	是

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        """
        Args:
            data: 原始数据
            labels: 标签
            transform: 可选的数据变换
        """
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]

        if self.transform:
            sample = self.transform(sample)

        return sample, label

Dataset 的设计使得 DataLoader 可以通过索引访问任意样本，从而支持随机打乱（shuffle）和多进程加载。

---

自定义 Dataset 实现

图片数据集

import os
import torch
from torch.utils.data import Dataset
from torchvision.io import read_image
from torchvision import transforms

class ImageDataset(Dataset):
    """自定义图片数据集，目录结构: root/class_name/image.jpg"""

    def __init__(self, root_dir, transform=None, target_transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.target_transform = target_transform

        # 收集所有图片路径和标签
        self.samples = []
        self.classes = sorted(os.listdir(root_dir))
        self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)}

        for cls_name in self.classes:
            cls_dir = os.path.join(root_dir, cls_name)
            if not os.path.isdir(cls_dir):
                continue
            for img_name in os.listdir(cls_dir):
                img_path = os.path.join(cls_dir, img_name)
                if img_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')):
                    self.samples.append((img_path, self.class_to_idx[cls_name]))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        img_path, label = self.samples[idx]

        # 使用 torchvision.io 读取图片（返回 tensor）
        image = read_image(img_path)

        # 转为 float 并归一化到 [0, 1]
        image = image.float() / 255.0

        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)

        return image, label

文本数据集

import torch
from torch.utils.data import Dataset

class TextDataset(Dataset):
    """自定义文本分类数据集"""

    def __init__(self, texts, labels, tokenizer, max_length=512):
        """
        Args:
            texts: 文本列表
            labels: 标签列表
            tokenizer: 分词器（如 HuggingFace tokenizer 或自定义）
            max_length: 最大序列长度
        """
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]

        # 分词并编码
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

        input_ids = encoding['input_ids'].squeeze(0)       # (max_length,)
        attention_mask = encoding['attention_mask'].squeeze(0)  # (max_length,)

        return {
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'label': torch.tensor(label, dtype=torch.long)
        }

CSV 数据集

import pandas as pd
import torch
import numpy as np
from torch.utils.data import Dataset

class CSVDataset(Dataset):
    """从 CSV 文件加载结构化数据"""

    def __init__(self, csv_path, feature_columns, label_column,
                 transform=None, dtype=torch.float32):
        """
        Args:
            csv_path: CSV 文件路径
            feature_columns: 特征列名列表
            label_column: 标签列名
            transform: 可选的特征变换函数
            dtype: 数据类型
        """
        df = pd.read_csv(csv_path)

        self.features = torch.tensor(df[feature_columns].values, dtype=dtype)
        self.labels = torch.tensor(df[label_column].values, dtype=torch.long)
        self.transform = transform

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        x = self.features[idx]
        y = self.labels[idx]

        if self.transform:
            x = self.transform(x)

        return x, y

---

内置数据集

torchvision.datasets 提供了大量常用数据集的封装，开箱即用。

常用内置数据集一览

数据集	说明	类型	自动下载
MNIST	手写数字 0-9，28x28 灰度图	图片分类	是
FashionMNIST	10 类时尚单品，28x28 灰度图	图片分类	是
CIFAR10	10 类彩色图片，32x32	图片分类	是
CIFAR100	100 类彩色图片，32x32	图片分类	是
ImageNet (ILSVRC2012)	1000 类大规模图片	图片分类	否（需申请）
ImageFolder	按文件夹组织的通用图片数据集	图片分类	-
CocoDetection	MS COCO 目标检测数据集	目标检测	否
VOCDetection	Pascal VOC 目标检测数据集	目标检测	是
CelebA	名人人脸属性数据集	图片属性	是
STL10	10 类高分辨率图片（96x96）	图片分类/自监督	是
SVHN	街景门牌号数据集	图片分类	是
Caltech101/256	Caltech 图片数据集	图片分类	是
Cityscapes	城市街景语义分割	语义分割	否

使用示例

import torchvision
from torchvision import transforms

# 定义预处理流水线
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# 下载并加载 CIFAR-10 训练集
train_dataset = torchvision.datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

# 下载并加载 MNIST 测试集
test_dataset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
)

# 使用 ImageFolder 加载自定义图片数据
# 目录结构: data/train/class1/xxx.jpg, data/train/class2/xxx.jpg
imagefolder_dataset = torchvision.datasets.ImageFolder(
    root='data/train',
    transform=transform
)

print(f"类别: {imagefolder_dataset.classes}")
print(f"类别到索引: {imagefolder_dataset.class_to_idx}")

ImageFolder 的目录结构要求

root/
├── class_1/
│   ├── img1.jpg
│   ├── img2.jpg
│   └── ...
├── class_2/
│   ├── img1.jpg
│   └── ...
└── class_N/
    └── ...

---

DataLoader 详解

DataLoader 是对 Dataset 的封装，提供批量加载、多进程读取、数据打乱等功能。

核心参数详解

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset=train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=True,
    collate_fn=None,
    worker_init_fn=None,
    persistent_workers=True,
    prefetch_factor=2,
    pin_memory_device="cuda"
)

参数说明表

参数	类型	默认值	说明
dataset	Dataset	必填	要加载的数据集
batch_size	int	1	每批样本数
shuffle	bool	False	是否在每个 epoch 打乱数据顺序
num_workers	int	0	数据加载的子进程数。0 表示在主进程中加载
pin_memory	bool	False	是否将数据拷贝到 CUDA 固定内存中，加速 GPU 传输
drop_last	bool	False	若最后一个 batch 样本数不足 batch_size，是否丢弃
collate_fn	callable	None	自定义样本合并函数
worker_init_fn	callable	None	每个 worker 进程的初始化函数
persistent_workers	bool	False	是否在 epoch 结束后保持 worker 进程存活
prefetch_factor	int	2	每个 worker 预取的 batch 数量
pin_memory_device	str	""	pin_memory 的目标设备
timeout	float	0	数据收集的超时时间（秒）。0 表示不超时
multiprocessing_context	str/Context	None	多进程上下文，如 'spawn'、'fork'、'forkserver'

典型训练循环

import torch
from torch.utils.data import DataLoader

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True if device.type == 'cuda' else False,
    persistent_workers=True
)

# 训练循环
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs = inputs.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if batch_idx % 100 == 0:
            print(f'Epoch [{epoch+1}], Step [{batch_idx}/{len(train_loader)}], '
                  f'Loss: {loss.item():.4f}')

    print(f'Epoch [{epoch+1}] 平均损失: {running_loss / len(train_loader):.4f}')

num_workers 调优建议

场景	建议 num_workers	说明
调试/小数据集	0	主进程加载，便于调试
一般训练	4-8	通常为 CPU 核心数的一半
大规模数据 + SSD	8-16	I/O 不是瓶颈时可增加
大规模数据 + HDD	4-8	机械硬盘随机读取慢，过多进程反而降低效率

提示：可以通过 num_workers=0 先确保 Dataset 实现正确，再逐步增加 worker 数量。

---

自定义 collate_fn

collate_fn 控制如何将单个样本合并为一个 batch。默认行为是将样本列表沿第 0 维堆叠。对于变长序列等场景，需要自定义。

处理变长序列与 Padding

import torch
from torch.nn.utils.rnn import pad_sequence

def collate_fn_padding(batch):
    """
    处理变长序列的 collate 函数
    batch: list of (input_ids, label) tuples
    """
    # 分离输入和标签
    input_ids_list = [item['input_ids'] for item in batch]
    attention_mask_list = [item['attention_mask'] for item in batch]
    labels = torch.stack([item['label'] for item in batch])

    # 对变长序列进行 padding（默认 pad 到 batch 中最长序列）
    input_ids_padded = pad_sequence(input_ids_list, batch_first=True, padding_value=0)
    attention_mask_padded = pad_sequence(attention_mask_list, batch_first=True, padding_value=0)

    return {
        'input_ids': input_ids_padded,
        'attention_mask': attention_mask_padded,
        'label': labels
    }


def collate_fn_image_with_bboxes(batch):
    """
    处理目标检测中不同图片有不同数量的 bounding box 的情况
    batch: list of (image, bboxes, labels) tuples
    """
    images = torch.stack([item[0] for item in batch])

    # 每张图片的 bbox 数量不同，记录每张图片的目标数
    num_boxes = [item[1].shape[0] for item in batch]

    # 将所有 bbox 和 labels 拼接在一起
    all_bboxes = torch.cat([item[1] for item in batch], dim=0)
    all_labels = torch.cat([item[2] for item in batch], dim=0)

    # 生成 batch index，标识每个 bbox 属于哪张图片
    batch_idx = torch.cat([
        torch.full((n,), i, dtype=torch.long) for i, n in enumerate(num_boxes)
    ])

    return images, all_bboxes, all_labels, batch_idx


# 使用
text_loader = DataLoader(
    text_dataset,
    batch_size=32,
    shuffle=True,
    collate_fn=collate_fn_padding,
    num_workers=2
)

---

torchvision.transforms 详解

torchvision.transforms 提供了丰富的图像变换操作，用于数据预处理和数据增强。

常用变换

from torchvision import transforms

# 基础变换组合
basic_transform = transforms.Compose([
    transforms.Resize(256),                 # 将短边缩放到 256
    transforms.CenterCrop(224),             # 从中心裁剪 224x224
    transforms.ToTensor(),                  # 转为 tensor，像素值归一化到 [0, 1]
    transforms.Normalize(                   # 标准化
        mean=[0.485, 0.456, 0.406],         # ImageNet 均值
        std=[0.229, 0.224, 0.225]           # ImageNet 标准差
    )
])

常用变换参数表

Transform	说明	关键参数
ToTensor()	PIL/numpy 转 tensor，像素 [0,255] -> [0.0, 1.0]	-
ToPILImage()	tensor/numpy 转 PIL Image	mode
Normalize(mean, std)	标准化: (x - mean) / std	mean, std
Resize(size)	缩放到指定大小	size, interpolation
CenterCrop(size)	中心裁剪	size
RandomCrop(size)	随机裁剪	size, padding
RandomResizedCrop(size)	随机裁剪并缩放	size, scale, ratio
Grayscale(num_output_channels)	转灰度图	num_output_channels
ConvertImageDtype(dtype)	转换图像数据类型	dtype
Lambda(func)	自定义 lambda 变换	func

数据增强

数据增强（Data Augmentation）通过对训练图片进行随机变换来扩充数据集，有效防止过拟合。

# 训练时使用数据增强
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.08, 1.0), ratio=(3/4, 4/3)),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(degrees=15),
    transforms.ColorJitter(
        brightness=0.4,
        contrast=0.4,
        saturation=0.4,
        hue=0.1
    ),
    transforms.RandomGrayscale(p=0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    transforms.RandomErasing(p=0.25, scale=(0.02, 0.33), ratio=(0.3, 3.3))
])

# 验证/测试时不做增强
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

数据增强 Transform 详解

Transform	说明	关键参数
RandomHorizontalFlip(p)	水平随机翻转	p: 翻转概率
RandomVerticalFlip(p)	垂直随机翻转	p: 翻转概率
RandomRotation(degrees)	随机旋转	degrees: 旋转角度范围
RandomAffine(degrees, translate, scale, shear)	随机仿射变换	多个参数
ColorJitter(b, c, s, h)	颜色抖动	亮度、对比度、饱和度、色调
RandomGrayscale(p)	随机转灰度	p: 转换概率
RandomPerspective(distortion_scale, p)	随机透视变换	失真比例和概率
GaussianBlur(kernel_size, sigma)	高斯模糊	核大小和 sigma 范围
RandomErasing(p, scale, ratio)	随机擦除（需在 ToTensor 后使用）	概率、面积比、纵横比
AutoAugment(policy)	自动增强策略	policy: ImageNet/CIFAR10/SVHN 等
RandAugment(num_ops, magnitude)	随机增强（RandAugment）	操作数量和强度
TrivialAugmentWide()	简单增强（TrivialAugment）	-

AutoAugment 使用

from torchvision import transforms

# 使用预定义的 AutoAugment 策略
auto_augment_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.AutoAugment(policy=transforms.AutoAugmentPolicy.IMAGENET),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# RandAugment
rand_augment_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandAugment(num_ops=2, magnitude=9),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

自定义 Transform

import torch
import numpy as np
from torchvision import transforms

class AddGaussianNoise:
    """添加高斯噪声的自定义 Transform"""

    def __init__(self, mean=0.0, std=0.1):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        noise = torch.randn_like(tensor) * self.std + self.mean
        return tensor + noise

    def __repr__(self):
        return f'{self.__class__.__name__}(mean={self.mean}, std={self.std})'


class Cutout:
    """Cutout: 随机遮挡图像中一个矩形区域"""

    def __init__(self, length):
        self.length = length

    def __call__(self, img):
        h, w = img.shape[1], img.shape[2]
        mask = torch.ones(h, w)

        y = np.random.randint(h)
        x = np.random.randint(w)

        y1 = np.clip(y - self.length // 2, 0, h)
        y2 = np.clip(y + self.length // 2, 0, h)
        x1 = np.clip(x - self.length // 2, 0, w)
        x2 = np.clip(x + self.length // 2, 0, w)

        mask[y1:y2, x1:x2] = 0.0
        mask = mask.unsqueeze(0).expand_as(img)
        return img * mask


# 组合使用
custom_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
    AddGaussianNoise(mean=0.0, std=0.05),
    # Cutout(16),  # 注意 Cutout 会影响后续操作
])

transforms.RandomChoice 与 transforms.RandomApply

# RandomApply: 以一定概率执行一组变换
transform = transforms.RandomApply([
    transforms.ColorJitter(brightness=0.4, contrast=0.4),
    transforms.RandomRotation(30)
], p=0.5)  # 50% 概率执行这组变换

# RandomChoice: 从一组变换中随机选择一个执行
transform = transforms.RandomChoice([
    transforms.RandomCrop(224),
    transforms.RandomResizedCrop(224),
    transforms.CenterCrop(224)
])

---

IterableDataset 适用场景与实现

当数据集太大无法全部载入内存，或数据来自流式源（如网络、日志、数据库流）时，应使用 IterableDataset。

IterableDataset vs Map-style Dataset

特性	Map-style Dataset	IterableDataset
继承	Dataset	IterableDataset
核心方法	__getitem__ + __len__	__iter__
随机访问	支持	不支持
适合数据量	中小规模	超大规模 / 流式
shuffle	DataLoader 自动支持	需要自行实现
多 worker	无特殊处理	需要处理 worker 分片

实现示例

import torch
from torch.utils.data import IterableDataset, DataLoader

class LargeTextDataset(IterableDataset):
    """大型文本文件的流式读取"""

    def __init__(self, file_path, tokenizer, max_length=512):
        self.file_path = file_path
        self.tokenizer = tokenizer
        self.max_length = max_length

    def parse_file(self, file_path):
        """逐行读取文件并生成样本"""
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                # 假设每行格式: "label\ttext"
                parts = line.split('\t', 1)
                if len(parts) != 2:
                    continue
                label_str, text = parts
                label = int(label_str)

                encoding = self.tokenizer(
                    text,
                    max_length=self.max_length,
                    padding='max_length',
                    truncation=True,
                    return_tensors='pt'
                )

                yield {
                    'input_ids': encoding['input_ids'].squeeze(0),
                    'attention_mask': encoding['attention_mask'].squeeze(0),
                    'label': torch.tensor(label, dtype=torch.long)
                }

    def __iter__(self):
        """处理多 worker 分片"""
        worker_info = torch.utils.data.get_worker_info()

        if worker_info is None:
            # 单 worker，直接迭代
            yield from self.parse_file(self.file_path)
        else:
            # 多 worker 时，将文件分成多段
            # 更精细的分片方式：每个 worker 读取不同文件
            total_size = self._get_file_size()
            per_worker = total_size // worker_info.num_workers
            start = worker_info.id * per_worker
            end = start + per_worker if worker_info.id != worker_info.num_workers - 1 else total_size

            yield from self._parse_file_range(self.file_path, start, end)

    def _get_file_size(self):
        """获取文件行数"""
        with open(self.file_path, 'r', encoding='utf-8') as f:
            return sum(1 for _ in f)

    def _parse_file_range(self, file_path, start, end):
        """读取文件指定行范围"""
        with open(file_path, 'r', encoding='utf-8') as f:
            for i, line in enumerate(f):
                if i < start:
                    continue
                if i >= end:
                    break
                line = line.strip()
                if not line:
                    continue
                parts = line.split('\t', 1)
                if len(parts) != 2:
                    continue
                label_str, text = parts
                label = int(label_str)

                encoding = self.tokenizer(
                    text,
                    max_length=self.max_length,
                    padding='max_length',
                    truncation=True,
                    return_tensors='pt'
                )

                yield {
                    'input_ids': encoding['input_ids'].squeeze(0),
                    'attention_mask': encoding['attention_mask'].squeeze(0),
                    'label': torch.tensor(label, dtype=torch.long)
                }


# 多文件 IterableDataset
class MultiFileIterableDataset(IterableDataset):
    """多个文件的流式数据集，自动分配给不同 worker"""

    def __init__(self, file_list, parse_fn):
        self.file_list = file_list
        self.parse_fn = parse_fn

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()

        if worker_info is None:
            files = self.file_list
        else:
            # 将文件列表均分给各个 worker
            per_worker = len(self.file_list) // worker_info.num_workers
            worker_id = worker_info.id
            start = worker_id * per_worker
            end = start + per_worker if worker_id != worker_info.num_workers - 1 else len(self.file_list)
            files = self.file_list[start:end]

        for file_path in files:
            yield from self.parse_fn(file_path)


# 使用
stream_dataset = LargeTextDataset('data/train.txt', tokenizer=my_tokenizer)
stream_loader = DataLoader(
    stream_dataset,
    batch_size=32,
    num_workers=4,
    # IterableDataset 不能设置 shuffle=True
)

for batch in stream_loader:
    # 训练...
    pass

---

多进程数据加载注意事项

常见陷阱

1. 随机种子问题：每个 worker 会有不同的随机种子，确保数据增强的随机性在 worker 间独立

def worker_init_fn(worker_id):
    """初始化每个 worker 的随机种子"""
    import numpy as np
    np.random.seed(np.random.get_state()[1][0] + worker_id)

dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    worker_init_fn=worker_init_fn
)

2. 共享资源问题：不要在 Dataset 中使用锁、文件句柄等不可序列化的共享资源

# 错误示范：在 __init__ 中打开文件句柄
class BadDataset(Dataset):
    def __init__(self, file_path):
        self.file = open(file_path, 'r')  # 这在多 worker 时会出问题！

# 正确做法：在 __getitem__ 中按需打开
class GoodDataset(Dataset):
    def __init__(self, file_path):
        self.file_path = file_path

    def __getitem__(self, idx):
        with open(self.file_path, 'r') as f:
            # 定位到第 idx 行并读取
            for i, line in enumerate(f):
                if i == idx:
                    return self._process(line)

3. CUDA 张量不能在 worker 进程中创建：数据应在主进程中移动到 GPU

# 正确做法
for batch in dataloader:
    batch = batch.to(device)  # 在主进程中移到 GPU

4. persistent_workers=True：避免每个 epoch 重新创建 worker 进程的开销，但会占用更多内存

fork vs spawn

import multiprocessing as mp

# Linux 默认使用 fork，macOS/Windows 默认使用 spawn
# fork: 子进程继承父进程的内存状态（快但可能有死锁问题）
# spawn: 子进程从头启动（安全但慢）

# 推荐在 Linux 上也使用 spawn 以避免 CUDA fork 问题
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    multiprocessing_context='spawn'
)

---

内存映射与大数据集处理技巧

使用 numpy 内存映射

import numpy as np
import torch
from torch.utils.data import Dataset

class MmapDataset(Dataset):
    """使用内存映射加载大型 numpy 数据文件"""

    def __init__(self, data_path, labels_path, transform=None):
        # mmap_mode='r' 不会将整个文件加载到内存
        self.data = np.load(data_path, mmap_mode='r')
        self.labels = np.load(labels_path, mmap_mode='r')
        self.transform = transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        # 访问时才从磁盘读取对应数据
        x = torch.from_numpy(self.data[idx].copy()).float()
        y = torch.tensor(self.labels[idx], dtype=torch.long)

        if self.transform:
            x = self.transform(x)

        return x, y


# 一次性保存大型数据集为 .npy 文件
def create_mmap_dataset(data, labels, data_path, labels_path):
    """将数据保存为 numpy 格式以支持内存映射"""
    np.save(data_path, data.astype(np.float32))
    np.save(labels_path, labels.astype(np.int64))

使用 HDF5

import h5py
import torch
import numpy as np
from torch.utils.data import Dataset

class HDF5Dataset(Dataset):
    """使用 HDF5 格式存储和读取大型数据集"""

    def __init__(self, hdf5_path, data_key='data', label_key='labels',
                 transform=None):
        self.hdf5_path = hdf5_path
        self.data_key = data_key
        self.label_key = label_key
        self.transform = transform

        # 获取数据集大小（不加载全部数据）
        with h5py.File(self.hdf5_path, 'r') as f:
            self.length = len(f[self.data_key])

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        # 每次访问时打开文件（兼容多 worker）
        with h5py.File(self.hdf5_path, 'r') as f:
            x = torch.from_numpy(f[self.data_key][idx].astype(np.float32))
            y = torch.tensor(f[self.label_key][idx], dtype=torch.long)

        if self.transform:
            x = self.transform(x)

        return x, y

使用 WebDataset 处理大规模数据

# WebDataset 支持从 tar 文件流式读取，适合超大规模数据集
# pip install webdataset
import webdataset as wds

# 将数据打包为 tar 文件: train-000000.tar, train-000001.tar, ...
dataset = (
    wds.WebDataset("path/to/train-{000000..000999}.tar")
    .shuffle(1000)
    .decode("pil")
    .to_tuple("jpg", "cls")
    .map_tuple(transform, lambda x: x)
    .batched(32)
)

常见优化技巧

技巧	说明
预处理并缓存	将耗时的预处理结果保存到磁盘，避免重复计算
内存映射	使用 mmap_mode='r' 避免一次性加载全部数据
合理 num_workers	根据 CPU 核心数和 I/O 类型调整
pin_memory + non_blocking	GPU 训练时开启，减少 CPU 到 GPU 的传输等待
persistent_workers	减少 worker 进程重建开销
预取（prefetch_factor）	提前加载后续 batch，隐藏 I/O 延迟
混合精度 + 高效格式	使用 FP16 存储可减少内存和 I/O

# 综合优化示例
import torch
from torch.utils.data import DataLoader

def create_optimized_dataloader(dataset, batch_size, num_workers=4, is_training=True):
    """创建优化后的 DataLoader"""
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=is_training,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        persistent_workers=num_workers > 0,
        prefetch_factor=2 if num_workers > 0 else None,
        drop_last=is_training,
        worker_init_fn=lambda worker_id: np.random.seed(
            np.random.get_state()[1][0] + worker_id
        )
    )

---

参考资料

• torch.utils.data 官方文档
• torchvision.datasets 官方文档
• torchvision.transforms 官方文档
• PyTorch Dataset Tutorial
• Data Loading and Processing Tutorial
• IterableDataset 官方文档