PyTorch数据加载机制详解：Dataset、DataLoader与采样策略优化
Dataset、DataLoader 与 Sampler 的协作机制
在 PyTorch 中，数据加载流程由三个核心组件构成：Dataset 负责数据存储与索引访问，Sampler 控制样本的选取顺序，而 DataLoader 则整合二者并实现批量加载。可以将整个过程类比为一个工厂流水线：
- Dataset：存放原始数据的仓库，支持通过索引（index）获取单个样本。
- Sampler：决定从仓库中按何种顺序提取索引，如顺序读取或随机打乱。
- DataLoader：雇佣多个工人（worker），根据 Sampler 提供的索引列表，从 Dataset 中并行读取数据，并组织成批次（batch）。
自定义 Dataset 实现方法
要创建可被 DataLoader 使用的数据集，需继承 torch.utils.data.Dataset 并实现以下三个方法：
__init__：初始化数据路径、预处理变换等。__getitem__：接收索引，返回对应的输入和标签。__len__：返回数据总量，用于控制迭代范围。
import torch
from torch.utils.data import Dataset
class CustomTensorDataset(Dataset):
def __init__(self):
self.samples = torch.arange(0, 20).float()
def __getitem__(self, idx):
x = self.samples[idx]
y = x * 2 + 1
return y
def __len__(self):
return len(self.samples)
使用示例：
# 初始化数据集
dataset = CustomTensorDataset()
print(len(dataset)) # 输出: 20
print(dataset[5]) # 输出: 11.0
# 构建数据加载器
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=4,
shuffle=True
)
print(len(loader)) # 输出: 5 (20/4)
for batch in loader:
print(batch)
关键参数解析：pin_memory 与 num_workers
pin_memory（页锁定内存）
当设置 pin_memory=True 时，DataLoader 会将张量存放在主机的页锁定内存（pinned memory）中。这种内存不会被交换到磁盘，因此能够显著加快从 CPU 内存向 GPU 显存的数据传输速度。
显存本质上都是页锁定的，若源内存也锁定，则可通过 DMA（直接内存访问）快速复制，无需阻塞 CPU。但该功能对物理内存要求较高，建议仅在 RAM 充足时启用。
num_workers（多进程加载）
此参数指定用于异步加载数据的子进程数量。每个 worker 进程独立执行数据读取和预处理任务，从而避免 I/O 操作阻塞主训练循环。
通常建议将其设为主机逻辑核心数，例如 8 核 CPU 可设为 num_workers=8。但过高数值可能导致内存争用或开销增加，需结合实际情况调整。
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=32,
shuffle=True,
num_workers=8,
pin_memory=True
)
DataLoader 核心参数说明
| 参数名 | 作用描述 |
|---|---|
dataset | 实现了 __getitem__ 和 __len__ 的数据集对象 |
batch_size | 每批返回的样本数量 |
shuffle | 是否在每个 epoch 开始前打乱数据顺序 |
sampler | 自定义索引生成策略，替代 shuffle |
batch_sampler | 以批为单位提供索引组 |
num_workers | 数据加载子进程数 |
collate_fn | 合并单个样本为 batch 的函数 |
pin_memory | 是否使用页锁定内存加速 GPU 传输 |
drop_last | 丢弃最后一个不完整 batch |
常用 Sampler 类型详解
All samplers inherit from the base Sampler class:
class Sampler:
def __init__(self, data_source):
self.data_source = data_source
def __iter__(self):
raise NotImplementedError
def __len__(self):
raise NotImplementedError
1. SequentialSampler
按自然顺序返回索引，适用于验证阶段。
class SequentialSampler(Sampler):
def __iter__(self):
return iter(range(len(self.data_source)))
2. RandomSampler
随机打乱索引序列。支持重复采样（with replacement）和限制采样数量。
sampler = RandomSampler(data_source, replacement=False, num_samples=1000)
3. SubsetRandomSampler
从指定索引子集中随机采样，常用于划分训练/验证集。
n_total = len(full_dataset)
n_val = n_total // 5
indices = list(range(n_total))
np.random.shuffle(indices)
train_idx = indices[n_val:]
val_idx = indices[:n_val]
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetRandomSampler(val_idx)
train_loader = DataLoader(dataset, sampler=train_sampler, batch_size=32)
4. WeightedRandomSampler
根据权重分布进行采样，适用于类别不平衡场景。
# 假设每个样本的采样权重
weights = [0.1, 0.5, 1.0, ...]
sampler = WeightedRandomSampler(weights, num_samples=1000, replacement=True)
5. BatchSampler
封装基础 sampler，将其输出的单个索引组合为 batch。
sampler = SequentialSampler(data_source)
batch_sampler = BatchSampler(sampler, batch_size=4, drop_last=False)
for batch_indices in batch_sampler:
print(batch_indices) # 如 [0,1,2,3], [4,5,6,7]
6. DistributedSampler
在分布式训练中确保各进程加载不同子集。
if distributed:
train_sampler = torch.utils.data.DistributedSampler(
dataset, num_replicas=world_size, rank=rank, shuffle=True
)
else:
train_sampler = None
data_loader = DataLoader(
dataset,
batch_size=16,
sampler=train_sampler,
num_workers=4,
pin_memory=True
)
