PyTorch DataLoader 在每个时期使用相同的随机变换

Posted 2023-02-16

【中文标题】PyTorch DataLoader 在每个时期使用相同的随机变换

【英文标题】：PyTorch DataLoader uses identical random transformation across each epoch

【发布时间】：2021-07-15 15:53:10

【问题描述】：

PyTorch/Numpy 中有一个bug，当与DataLoader 并行加载批次时（即设置num_workers > 1），每个worker 使用相同的NumPy 随机种子，导致应用的任何随机函数为跨并行批次相同。这可以通过将种子生成器传递给worker_init_fn 参数like so 来解决。

但问题在多个时期仍然存在。

小例子：

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 4

dataset = RandomDataset()
dataloader = DataLoader(dataset, batch_size=1, 
                        num_workers=2, 
                        worker_init_fn = lambda x: np.random.seed(x))

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

如您所见，虽然一个时期内的并行批次现在产生不同的结果，但跨时期的结果是相同的：

Epoch 0
tensor([[684, 559]])
tensor([[ 37, 235]])
tensor([[629, 192]])
tensor([[908,  72]])

Epoch 1
tensor([[684, 559]])
tensor([[ 37, 235]])
tensor([[629, 192]])
tensor([[908,  72]])

Epoch 2
tensor([[684, 559]])
tensor([[ 37, 235]])
tensor([[629, 192]])
tensor([[908,  72]])

如何解决这种行为？

---

使用空参数，例如worker_init_fn = lambda _: np.random.seed() 似乎解决了这个问题 - 这个解决方法有什么问题吗？

【参考方案1】：

我能想到的最好的方法是使用pytorch设置的种子来进行numpy和random：

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, 1000, 2)

    def __len__(self):
        return 4

dataset = RandomDataset()
dataloader = DataLoader(dataset, batch_size=1, 
                        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 秒。

【参考方案2】：

正如您链接的博客文章中所述，您所写的内容将在每个时期为每个工人产生相同的随机数：

在数据集上迭代 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

【讨论】：

您能否详细说明一下如何您使用“火炬随机”来实现这一目标？