PyTorch迁移至Datapipe后如何实现Subset截取部分数据功能

旧版Dataset体系下用torch.utils.data.Subset取子集的逻辑，在DataPipe体系下有原生的对应实现，不需要自定义封装，按使用场景选对应方法即可：

• 场景1：截取前N条连续数据（例如前1000条）
  这是最常用的场景，直接调用DataPipe内置的header()方法即可，性能最优，不会加载全量数据再做截断：

  from torchdata.datapipes.iter import IterableWrapper
  
  # 替换为你实际使用的DataPipe数据源
  raw_dp = IterableWrapper(range(100000))
  # 截取前1000条，等价于旧版Subset取索引0-999的效果
  subset_dp = raw_dp.header(1000)
  

  注意：header()对Iter-style、Map-style两类DataPipe都适用，会自动在读取到指定条数后终止遍历，不会额外消耗内存。

• 场景2：按指定索引列表取非连续子集
  如果你需要和Subset(indices=...)完全一致的能力，按自定义索引列表取数，使用Map-style DataPipe内置的subset()方法：

  # 自定义要取用的索引列表
  selected_indices = [0, 3, 7, 10] + list(range(200, 500))
  # 注意：subset仅支持实现了随机索引访问的Map-style DataPipe
  subset_dp = map_style_dp.subset(selected_indices)
  

  如果你用的是仅支持顺序遍历的Iter-style DataPipe，需要按索引筛选，可以先给数据加索引再过滤：

  # 先为顺序遍历的数据绑定索引
  dp_with_idx = raw_dp.enumerate()
  # 按索引条件过滤后去掉索引字段，得到目标子集
  subset_dp = dp_with_idx.filter(
      lambda x: x[0] < 1000  # 这里替换成你需要的索引判断逻辑
  ).map(lambda x: x[1])
  

注意事项：不要通过手动写循环、触发break的方式截断数据，这种写法会破坏DataPipe的执行图优化，也无法兼容DataLoader的多进程加载、预取等流水线能力，优先使用内置的截取方法。

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