PyTorch DataLoader 对并行运行的批次使用相同的随机种子

Question

在 PyTorch/Numpy 中有一个 bug，当与 DataLoader 并行加载批次时（即设置 num_workers > 1），每个工作人员使用相同的 NumPy 随机种子，导致应用的任何随机函数都是在并行批次中相同。

小例子：

import numpy as np
from torch.utils.data import Dataset, DataLoader

class RandomDataset(Dataset):
    def __getitem__(self, index):
        return np.random.randint(0, 1000, 2)

    def __len__(self):
        return 9
    
dataset = RandomDataset()
dataloader = DataLoader(dataset, batch_size=1, num_workers=3)

for batch in dataloader:
    print(batch)

如您所见，对于每个并行化的批次集 (3)，结果是相同的：

# First 3 batches
tensor([[891, 674]])
tensor([[891, 674]])
tensor([[891, 674]])
# Second 3 batches
tensor([[545, 977]])
tensor([[545, 977]])
tensor([[545, 977]])
# Third 3 batches
tensor([[880, 688]])
tensor([[880, 688]])
tensor([[880, 688]])

解决此问题的推荐/最优雅的方法是什么？即让每批产生不同的随机化，而不管工人的数量。

Answer 1

这似乎可行，至少在 Colab 中是这样：

dataloader = DataLoader(dataset, batch_size=1, num_workers=3, 
    worker_init_fn = lambda id: np.random.seed(id) )

编辑：

当迭代多个时期时，它会产生相同的输出（即相同的问题）。 – iacob

到目前为止我发现的最佳解决方案：

...
dataloader = DataLoader(ds, num_workers=2, 
           worker_init_fn = lambda id: np.random.seed(id + epoch * 10 ))

for epoch in range ( 2 ):
    for batch in dataloader:
        print(batch)
    print()

仍然不能建议封闭形式，事情取决于然后调用的 var (epoch)。理想情况下，它必须类似于worker_init_fn = lambda id: np.random.seed(id + EAGER_EVAL(np.random.randint(10000) )，其中 EAGER_EVAL 在将 lambda 作为参数传递之前评估加载程序构造的种子。我想知道在python中是否可能。