引言
在深度学习模型训练中，torch.utils.data.DataLoader是PyTorch提供的一个核心工具，用于高效地加载数据。通常，我们会为其指定一个固定的batch_size参数，使得每个训练批次都包含相同数量的样本。然而，在某些高级或特定场景下，我们可能需要更灵活的批处理策略，例如，根据数据样本的特性（如长度、复杂性）或硬件内存限制，动态地调整每个批次的样本数量。例如，我们可能希望在训练的不同阶段或处理不同类型的数据时，使用一系列预设的批大小[30, 60, 110, ..., 231]，而不是单一的64。
PyTorch的DataLoader通过其sampler和batch_sampler参数提供了极大的灵活性，允许用户自定义数据样本的索引生成逻辑。本文将详细介绍如何通过实现自定义的BatchSampler来满足动态批处理的需求。
PyTorch DataLoader与批处理机制
DataLoader的核心功能是迭代地从Dataset中获取数据批次。其工作流程大致如下：
- Dataset负责存储和按索引获取单个样本。
- Sampler（或默认的SequentialSampler/RandomSampler）负责生成单个样本的索引序列。
- BatchSampler（或默认的BatchSampler，它基于Sampler和batch_size生成批次索引列表）负责将这些单个样本索引组合成批次索引列表。
- DataLoader接收这些批次索引列表，从Dataset中取出对应的样本，并通过collate_fn将它们组合成张量批次。
当我们使用batch_size参数时，DataLoader内部会默认创建一个BatchSampler来按照固定大小对索引进行批处理。要实现动态批处理，我们需要绕过这个默认行为，提供一个能够生成可变大小批次索引的自定义BatchSampler。
实现自定义动态批次采样器（VariableBatchSampler）
为了实现动态批处理，我们将创建一个继承自torch.utils.data.Sampler的自定义类VariableBatchSampler。尽管其名称为Sampler，但其内部逻辑是直接生成批次索引，使其更适合作为DataLoader的batch_sampler参数使用。
import torch
from torch.utils.data import Sampler, TensorDataset, DataLoader
class VariableBatchSampler(Sampler):
"""
一个自定义的批次采样器，根据预定义的批大小列表生成可变大小的批次索引。
"""
def __init__(self, dataset_len: int, batch_sizes: list):
"""
初始化VariableBatchSampler。
Args:
dataset_len (int): 数据集的总长度（样本数量）。
batch_sizes (list): 一个包含每个批次所需样本数量的列表。
列表中所有元素的和应等于或大于dataset_len。
"""
if not isinstance(batch_sizes, list) or not all(isinstance(bs, int) and bs > 0 for bs in batch_sizes):
raise ValueError("batch_sizes 必须是一个包含正整数的列表。")
if sum(batch_sizes) < dataset_len:
print(f"警告: 提供的批大小总和 ({sum(batch_sizes)}) 小于数据集长度 ({dataset_len})。部分数据可能不会被采样。")
self.dataset_len = dataset_len
self.batch_sizes = batch_sizes
self.batch_idx = 0 # 当前正在处理的批次大小在batch_sizes列表中的索引
self.start_idx = 0 # 当前批次的起始索引
# 初始设置当前批次的结束索引。
# 如果batch_sizes列表为空，则默认为0，但前置检查会避免这种情况。
self.end_idx = self.batch_sizes[self.batch_idx] if self.batch_sizes else 0
def __iter__(self):
"""
返回采样器自身，使其可迭代。
在每次新的迭代开始时，重置状态。
"""
self.batch_idx = 0
self.start_idx = 0
self.end_idx = self.batch_sizes[self.batch_idx] if self.batch_sizes else 0
return self
def __next__(self):
"""
生成下一个批次的索引。
"""
if self.start_idx >= self.dataset_len:
# 如果起始索引已超出数据集长度，则表示所有数据已采样完毕
raise StopIteration()
# 获取当前批次的索引范围
# 注意：这里的索引是顺序生成的。如果需要随机批次，需要先打乱整个数据集的索引。
batch_indices = torch.arange(self.start_idx, min(self.end_idx, self.dataset_len), dtype=torch.int64)
# 更新起始索引为当前批次的结束位置
self.start_idx = min(self.end_idx, self.dataset_len)
self.batch_idx += 1 # 移动到下一个批次大小
# 尝试更新下一个批次的结束索引
try:
self.end_idx += self.batch_sizes[self.batch_idx]
except IndexError:
# 如果batch_sizes列表已用尽，将结束索引设置为数据集的末尾，
# 确保最后一个批次包含所有剩余的样本
self.end_idx = self.dataset_len
return batch_indices
VariableBatchSampler解析
-
__init__(self, dataset_len: int, batch_sizes: list):
- dataset_len: 数据集的总样本数。
- batch_sizes: 一个列表，其中每个元素代表一个批次的大小。这个列表的顺序决定了批次生成的顺序。重要提示：此列表中所有批次大小的总和应等于数据集的总长度，以确保所有数据都被采样且没有重复。
- self.batch_idx: 用于追踪当前正在使用batch_sizes列表中哪个批次大小。
- self.start_idx: 当前批次的起始索引。
- self.end_idx: 当前批次的结束索引（不包含）。
-
__iter__(self):
- 使采样器对象可迭代。每次新的迭代开始时（例如，每个epoch开始时），会重置batch_idx、start_idx和end_idx，确保从头开始采样。
-
__next__(self):
- 这是生成每个批次索引的核心逻辑。
- 首先检查self.start_idx是否已达到或超过self.dataset_len，如果是，则表示所有数据已采样完毕，抛出StopIteration。
- batch_indices = torch.arange(self.start_idx, min(self.end_idx, self.dataset_len), dtype=torch.int64)：生成从start_idx到end_idx（不包含）的索引张量。min(self.end_idx, self.dataset_len)确保不会超出数据集的实际范围，这对于处理最后一个批次可能比预设batch_size小的情况尤其重要。
- self.start_idx = min(self.end_idx, self.dataset_len)：更新下一个批次的起始索引。
- self.batch_idx += 1：移动到batch_sizes列表中的下一个批次大小。
- try-except IndexError块：尝试根据下一个批次大小更新self.end_idx。如果batch_sizes列表已耗尽（IndexError），则将self.end_idx设置为self.dataset_len，确保最后一个批次能够包含所有剩余的样本。
与DataLoader集成
VariableBatchSampler设计为直接作为DataLoader的batch_sampler参数。当使用batch_sampler时，DataLoader会期望它直接返回一个包含批次索引的列表或张量，并且DataLoader自身的batch_size参数会被忽略。
# 示例数据
x_train = torch.randn(8400, 4) # 8400个样本，每个样本4个特征
y_train = torch.randint(0, 2, (8400,)) # 8400个标签
train_dataset = TensorDataset(x_train, y_train)
# 定义动态批大小列表
# 确保所有批大小的总和等于数据集长度
list_batch_size = [30, 60, 110] * 20 + [8400 - sum([30, 60, 110] * 20)] # 示例：总和为8400
# 验证总和
assert sum(list_batch_size) == len(train_dataset), "批大小列表的总和必须等于数据集长度"
# 实例化自定义批次采样器
variable_batch_sampler = VariableBatchSampler(
dataset_len=len(train_dataset),
batch_sizes=list_batch_size
)
# 使用自定义批次采样器实例化DataLoader
# 注意：当使用batch_sampler时，batch_size参数会被忽略
data_loader_dynamic = DataLoader(
train_dataset,
batch_sampler=variable_batch_sampler,
num_workers=0 # 示例中设置为0，实际应用可根据需要设置
)
# 迭代DataLoader并打印每个批次的形状
print(f"数据集总样本数: {len(train_dataset)}")
print(f"动态批大小列表: {list_batch_size[:5]}... (共 {len(list_batch_size)} 个批次)")
for i, (data, labels) in enumerate(data_loader_dynamic):
print(f"批次 {i+1}: 数据形状 {data.shape}, 标签形状 {labels.shape}")
# 验证批次大小是否与预期一致
expected_batch_size = list_batch_size[i]
if i == len(list_batch_size) - 1 and sum(list_batch_size[:-1]) < len(train_dataset):
# 最后一个批次可能包含所有剩余样本，不一定严格等于list_batch_size的最后一个元素
# 除非list_batch_size精确求和等于dataset_len
pass # 这里的assert需要更复杂的逻辑，暂时跳过严格相等检查
else:
assert data.shape[0] == expected_batch_size, f"批次 {i+1} 大小不匹配。预期: {expected_batch_size}, 实际: {data.shape[0]}"
if i >= 10: # 仅打印前10个批次作为示例
print("...")
break
print("\n所有批次迭代完毕。")重要提示：
- 当将VariableBatchSampler作为batch_sampler参数传递给DataLoader时，DataLoader的batch_size参数应被省略或设置为默认值（1），因为它将被batch_sampler的逻辑覆盖。
- 如果将VariableBatchSampler作为sampler参数传递，DataLoader会默认batch_size=1，导致每个迭代返回的张量会多一个维度（例如，[batch_size, 1, features]），这通常不是我们想要的。因此，强烈建议使用batch_sampler。
注意事项与扩展
- 批大小总和与数据集长度：确保batch_sizes列表中所有元素的总和等于dataset_len。如果总和小于dataset_len，部分数据将不会被采样；如果总和大于dataset_len，__next__方法中的min(self.end_idx, self.dataset_len)会确保不会尝试采样超出数据集范围的索引，但可能会导致最后一个批次比list_batch_size中预期的要小。
- 随机性：上述VariableBatchSampler是顺序生成批次的。这意味着它总是从数据集的开头开始，并按照batch_sizes的顺序依次取出批次。如果需要随机的动态批次，您需要在__iter__方法中首先生成一个打乱的全局索引序列（例如，torch.randperm(self.dataset_len)），然后__next__方法从这个打乱的序列中按照batch_sizes指定的数量进行切片。
- drop_last行为：使用自定义BatchSampler时，DataLoader的drop_last参数不再直接生效，因为批次的生成完全由BatchSampler控制。如果您需要类似drop_last的功能（即丢弃最后一个不完整的批次），您需要在VariableBatchSampler的逻辑中自行实现。当前实现会尽可能地包含所有数据，即使最后一个批次小于预期的batch_size。
- 多进程数据加载：当使用num_workers > 0进行多进程数据加载时，BatchSampler的实例会在每个worker进程中被克隆。确保您的BatchSampler在多进程环境下能够正确工作，例如，如果它内部维护了复杂的状态，需要考虑如何同步或独立初始化这些状态。对于本教程中的VariableBatchSampler，由于其状态（batch_idx, start_idx, end_idx）在每个__iter__调用时都会重置，因此通常不会有大的问题。
总结
通过实现自定义的VariableBatchSampler，我们成功地为PyTorch的DataLoader引入了动态批处理的能力。这种方法提供了极高的灵活性，允许开发者根据特定的训练需求或数据特性，精确控制每个批次的数据量。无论是为了优化内存使用、处理变长序列，还是实现复杂的训练策略，自定义BatchSampler都是一个强大而专业的工具，能够显著提升数据加载和模型训练的效率与适应性。掌握这一技术，将使您在PyTorch深度学习开发中拥有更强的控制力。
