PyTorch DataLoader 在每个 epoch 中使用相同的随机变换 附加说明

如您链接的博文所述，您所写的内容将为每个时代的每个工人生成相同的随机数：

在数据集上迭代 3 次会在每个 epoch 产生相同的随机数。发生这种情况是因为对随机状态的所有更改对于每个工人都是本地的。默认情况下，工作进程在每个 epoch 结束时被杀死，并且所有工作资源都将丢失。同时，主进程中的随机状态没有改变，用于再次初始化各个工作进程。

给出解决方案：

因此您需要在每个 epoch 更改 NumPy 的种子，例如通过 np.random.seed(initial_seed + epoch)。

但我个人更喜欢只使用 Torch random 而不是 Numpy 的，以避免出现问题，因为 Torch 默认处理并行代码中的随机性。

附加说明

根据博文：

PyTorch 通过自动将 [...] 种子设置为 seed + worker_id 来处理这些问题。

这意味着在您的数据集类或训练循环中使用 Pytorch random 函数不应跨批次或时期复制随机性。例如，您编写的最小示例可以像这样修复：

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

class RandomDataset(Dataset):
    def __getitem__(self,indeX):
        return torch.randint(0,1000,(2,))

    def __len__(self):
        return 4

dataset = RandomDataset()
dataloader = DataLoader(dataset,batch_size=1,num_workers=2)

for epoch in range(3):
    print(f'\nEpoch {epoch}')
    for batch in dataloader:
        print(batch)

大多数时候 Numpy 可以用 Torch random（或 Python random）代替。这是另一个用于图像分割的随机变换示例：

class RandomHorizontalFlip:

    def __init__(self,prob=0.5):
        self.prob = prob

    def __call__(self,input,target):
        if torch.randn(1).item() < self.prob:
            return F.hflip(input),F.hflip(target)
        else:
            return input,target

我能想到的最好的方法是使用pytorch为numpy和random设置的seed：

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

def worker_init_fn(worker_id):
    torch_seed = torch.initial_seed()
    random.seed(torch_seed + worker_id)
    if torch_seed >= 2**30:  # make sure torch_seed + workder_id < 2**32
        torch_seed = torch_seed % 2**30
    np.random.seed(torch_seed + worker_id)

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

    def __len__(self):
        return 4

dataset = RandomDataset()
dataloader = DataLoader(dataset,num_workers=2,worker_init_fn = worker_init_fn)

for epoch in range(3):
    print(f'\nEpoch {epoch}')
    for batch in dataloader:
        print(batch)

输出：

Epoch 0
tensor([[593,191]])
tensor([[207,469]])
tensor([[976,714]])
tensor([[ 13,119]])

Epoch 1
tensor([[836,664]])
tensor([[138,836]])
tensor([[409,313]])
tensor([[  2,221]])

Epoch 2
tensor([[269,888]])
tensor([[315,619]])
tensor([[892,774]])
tensor([[ 70,771]])

或者，您可以使用 int(time.time()) 来播种 numpy 和 random，假设每个 epoch 的运行时间超过 1 秒。