PyTorch DataLoader动态批处理：实现可变批大小训练

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

# 示例数据
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所有批次迭代完毕。")