如何基于已有PyTorch DataLoader创建子集DataLoader（不使用Subset）

下面提供几种符合需求的实现方式，无需直接对原Dataset使用torch.utils.data.Subset：

方法1：收集样本后构建新DataLoader #

这种方法会从原DataLoader中抽取指定数量的样本，存储到内存后创建新的Dataset和DataLoader，适合样本量较小的场景。

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

def create_subset_data_loader(data_loader, subset_size):
    samples = []
    targets = []
    remaining = subset_size
    
    # 从原DataLoader中抽取样本
    for batch_x, batch_y in data_loader:
        if remaining <= 0:
            break
        # 取当前batch中需要的样本数
        take = min(remaining, batch_x.shape[0])
        samples.append(batch_x[:take])
        targets.append(batch_y[:take])
        remaining -= take
    
    if remaining > 0:
        raise ValueError(f"原DataLoader仅包含 {subset_size - remaining} 个样本，不足请求的 {subset_size} 个")
    
    # 拼接样本张量并创建新Dataset
    samples_tensor = torch.cat(samples, dim=0)
    targets_tensor = torch.cat(targets, dim=0)
    subset_dataset = TensorDataset(samples_tensor, targets_tensor)
    
    # 复用原DataLoader的配置参数
    return DataLoader(
        subset_dataset,
        batch_size=data_loader.batch_size,
        shuffle=data_loader.shuffle,
        num_workers=data_loader.num_workers,
        pin_memory=data_loader.pin_memory
    )

使用示例：

# 假设data_loader_original是已有的DataLoader实例
data_loader_subset = create_subset_data_loader(data_loader_original, 500)

方法2：实现切片语法（猴子补丁） #

如果想要像列表一样用data_loader_original[:size]的方式获取子集，可以给DataLoader类添加__getitem__方法。注意：这种方式会修改PyTorch内置类的行为，生产环境需谨慎使用。

import torch
from torch.utils.data import DataLoader

# 先定义方法1中的create_subset_data_loader函数（同上）

def dataloader_getitem(self, key):
    if isinstance(key, slice):
        start = key.start or 0
        stop = key.stop
        if stop is None:
            raise ValueError("切片必须指定结束位置")
        return create_subset_data_loader(self, stop - start)
    elif isinstance(key, int):
        # 返回单个样本（从对应batch中提取第一个元素）
        for idx, (batch_x, batch_y) in enumerate(self):
            if idx == key:
                return batch_x[0], batch_y[0]
        raise IndexError("索引超出范围")
    else:
        raise TypeError("仅支持整数索引或切片")

# 给DataLoader类添加__getitem__方法
DataLoader.__getitem__ = dataloader_getitem

使用示例：

# 获取前300个样本的子集DataLoader
data_loader_subset = data_loader_original[:300]

注意事项 #

• 若原DataLoader开启了shuffle=True，每次创建子集时抽取的样本会不同，如需固定子集，需先关闭原DataLoader的洗牌或设置固定随机种子。
• 当子集规模较大时，内存占用会显著增加，需根据硬件资源调整。

内容的提问来源于stack exchange，提问作者user3731622