PyTorch DataLoader Sampler自定义采样策略
在深度学习项目中，我们常常遇到这样的问题：模型训练初期损失下降缓慢，准确率停滞不前，尤其是当数据集存在严重类别不平衡时——比如医学影像中罕见病样本仅占5%。你可能会尝试调整学习率、换用更复杂的网络结构，但效果依然有限。其实，问题的根源可能并不在模型本身，而在于数据是如何被“喂”给模型的。
PyTorch 的DataLoader看似只是一个简单的数据加载工具，但它背后隐藏着一个关键机制：Sampler。正是这个组件决定了每一轮训练中，哪些样本会被选中、以什么顺序进入模型。标准的随机采样在面对复杂任务时往往力不从心，而通过自定义Sampler，我们可以精准调度数据流，让模型“看得更清楚”，从而显著提升训练效率和最终性能。
DataLoader是 PyTorch 中用于批量读取数据的核心类，它封装了数据集并提供高效的迭代接口。其核心参数之一sampler决定了样本索引的生成方式。所有采样器都继承自抽象基类torch.utils.data.Sampler，只需实现__iter__()方法即可定义自己的采样逻辑。
常见的内置采样器包括：
-SequentialSampler：按顺序遍历；
-RandomSampler：随机打乱；
-WeightedRandomSampler：根据权重概率抽样；
-BatchSampler：将多个索引组合成 batch。
但当我们需要处理如类别均衡、难样本挖掘或 episode-based 学习等场景时，就必须动手写一个自定义Sampler。
下面是一个典型的类别平衡采样器实现：
from torch.utils.data import Sampler, Dataset import torch import numpy as np class BalancedClassSampler(Sampler): """ 强制每个类别在每轮训练中被均匀采样的采样器。 假设 dataset.targets 包含标签列表。 """ def __init__(self, dataset: Dataset, num_samples_per_epoch: int): self.dataset = dataset self.num_samples_per_epoch = num_samples_per_epoch if not hasattr(dataset, 'targets'): raise ValueError("Dataset must have 'targets' attribute") self.targets = np.array(dataset.targets) self.classes = np.unique(self.targets) self.num_classes = len(self.classes) # 构建每个类别的索引池 self.class_indices = { cls: np.where(self.targets == cls)[0].tolist() for cls in self.classes } self.samples_per_class = self.num_samples_per_epoch // self.num_classes def __iter__(self): indices = [] # 主循环：从每个类别轮流采样 for _ in range(self.samples_per_class): for cls in self.classes: idx_list = self.class_indices[cls] random_idx = np.random.choice(idx_list) indices.append(random_idx) # 补足因整除导致的数量缺口 remaining = self.num_samples_per_epoch - len(indices) for _ in range(remaining): cls = np.random.choice(self.classes) random_idx = np.random.choice(self.class_indices[cls]) indices.append(random_idx) # 最终打乱顺序，避免类别间出现固定模式 np.random.shuffle(indices) return iter(indices) def __len__(self): return self.num_samples_per_epoch这个采样器的关键设计思想是：不让多数类垄断训练过程。例如在一个三分类任务中，若正常样本占90%，病变样本仅占10%，使用默认随机采样会导致模型很少“看到”病变案例。而通过强制每个类别等频出现，模型被迫关注少数类，有助于提高召回率和F1分数。
使用也非常简单：
dataset = MyDataset(...) sampler = BalancedClassSampler(dataset, num_samples_per_epoch=1200) dataloader = DataLoader(dataset, batch_size=32, sampler=sampler, num_workers=4)⚠️ 注意：一旦指定了
sampler参数，就不能再设置shuffle=True，否则会抛出异常。
这套机制的强大之处在于它的解耦性——Sampler与Dataset完全独立。这意味着你可以为同一个数据集搭配不同的采样策略，无需修改任何数据加载逻辑。更重要的是，在 GPU 加速环境下，这种灵活性能够快速转化为实际收益。
目前主流开发环境多基于容器化部署，例如预装 PyTorch 2.9 和 CUDA 12.1 的镜像pytorch-cuda:v2.9。这类镜像通常基于 NVIDIA 官方基础镜像构建，集成 Python 3.9+、cuDNN、Jupyter Notebook 及 SSH 服务，开箱即用。
启动后可直接验证 GPU 可用性：
import torch print("CUDA Available:", torch.cuda.is_available()) # 应返回 True print("GPU Count:", torch.cuda.device_count()) device = torch.device("cuda" if torch.cuda.is_available() else "cpu")在这种环境中运行前述采样器，配合多进程加载（num_workers=4），可以轻松实现高吞吐的数据供给。以下是一个端到端训练示例：
# 模拟一个不平衡数据集 class ImbalancedDataset(Dataset): def __init__(self): self.data = np.random.rand(1000, 3, 224, 224).astype(np.float32) labels_0 = [0] * 800 labels_1 = [1] * 150 labels_2 = [2] * 50 self.targets = labels_0 + labels_1 + labels_2 def __len__(self): return len(self.targets) def __getitem__(self, index): return torch.tensor(self.data[index]), self.targets[index] # 训练流程 dataset = ImbalancedDataset() sampler = BalancedClassSampler(dataset, num_samples_per_epoch=600) dataloader = DataLoader(dataset, batch_size=32, sampler=sampler, num_workers=4) model = torch.nn.Linear(3*224*224, 3).to(device) optimizer = torch.optim.Adam(model.parameters()) criterion = torch.nn.CrossEntropyLoss() for epoch in range(3): model.train() for x, y in dataloader: x = x.view(x.size(0), -1).to(device) y = y.to(device) optimizer.zero_grad() output = model(x) loss = criterion(output, y) loss.backward() optimizer.step() print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")整个流程无缝衔接，无需关心底层依赖。尤其是在服务器集群中，借助 Docker 挂载数据卷（-v /data:/workspace/data）和共享内存优化（--shm-size="8gb"），可以稳定支撑大规模训练任务。
实际应用中，自定义采样策略的价值远不止于类别平衡。比如在推荐系统中，用户行为稀疏，大多数商品从未被点击。如果只是随机采样负样本，模型很难学会区分“真正感兴趣的”和“只是没看过”的物品。这时就可以设计一个HardNegativeSampler，结合用户历史行为，优先选择那些语义相近但未被交互的“难负样本”。
类似地，在对比学习中，正负样本的构造直接影响表示质量；在元学习中，episode sampling 要求每次从不同任务中抽取支持集和查询集——这些高级训练范式都依赖于精细控制的采样逻辑。
然而，在编写自定义Sampler时也有一些工程细节必须注意：
| 实践建议 | 说明 |
|---|---|
| 线程安全 | 当num_workers > 0时，__iter__()可能在多个子进程中并发调用。避免使用全局状态或共享变量，确保每次返回的是全新迭代器。 |
| 长度一致性 | __len__()应尽量准确反映每个 epoch 的样本数，便于进度条显示和学习率调度器工作。 |
| 内存效率 | 不要在__init__中加载原始数据，只保存索引映射即可。对于超大数据集，甚至可以考虑懒加载索引文件。 |
| 分布式兼容性 | 在 DDP 训练中，需将自定义采样器包装在DistributedSampler中，或自行实现 rank 分片逻辑，防止不同 GPU 处理重复数据。 |
| 可复现性 | 若需实验结果可复现，应在__iter__中设置随机种子，例如np.random.seed(epoch + self.rank)。 |
此外，合理配置num_workers也很关键——一般建议不超过 CPU 核心数的70%，以免造成过多进程竞争资源。
从系统架构角度看，Dataset、Sampler和DataLoader共同构成了数据供给的三层体系：
+------------------+ +--------------------+ | Dataset |<----->| Custom Sampler | +------------------+ +--------------------+ ↓ +------------------+ | DataLoader | ——> 提供 batch 数据 +------------------+ ↓ +------------------+ | Model Training | ——> 在 GPU 上执行前向/反向传播 +------------------+ ↑ +------------------+ | PyTorch-CUDA Env | | (via Docker Img) | +------------------+这一设计体现了良好的职责分离：数据存储、访问策略、批处理逻辑各司其职。结合容器化环境提供的标准化运行平台，研究人员得以将精力集中在算法创新而非环境调试上。
当你下次面对模型表现不佳的问题时，不妨先问问自己：是不是数据的“出场顺序”出了问题？也许答案不在模型结构里，而在那个不起眼的Sampler中。掌握这项技能，意味着你不仅会“训练模型”，更能“指挥数据”——而这，正是现代深度学习工程能力的核心体现之一。
