如何在 pytorch 中重置数据加载器？

Question

我试图手动重置数据加载器，但无法。我在这里尝试了一切https://discuss.pytorch.org/t/how-could-i-reset-dataloader-or-count-data-batch-with-iter-instead-of-epoch/22902/4但没有运气. 任何人都知道如何重置数据加载器，并且批次的随机性/随机性不会被破坏？

Answer 1

要重置DataLoader，只需再次枚举加载器。每次调用都enumerate(loader)从头开始。

为了不破坏使用随机值的转换器，请在每次初始化 DataLoader 时重置随机种子。

def seed_init_fn(x):
   seed = args.seed + x
   np.random.seed(seed)
   random.seed(seed)
   torch.manual_seed(seed)
   return

loader = torch.utils.data.DataLoader(...., worker_init_fn = seed_init_fn)

while True:
   for i,data in enumerate(loader):
      # will always yield same data

见worker_init_fn文档：

https://pytorch.org/docs/master/data.html#torch.utils.data.DataLoader

这是一个更好的例子：

https://github.com/pytorch/pytorch/issues/5059#issuecomment-404232359