Seed Everything - 可复现的 PyTorch（一）

该系列其他文章

确定性算法 - 可复现的 PyTorch（二）

为了保证实验的「可复现性」，许多机器学习的代码都会有一个方法叫 seed_everything，这个方法尝试固定随机种子以让一些随机的过程在每一次的运行中产生相同的结果。但如果用谷歌搜索「how to seed everything in pytorch」，会得到各种不同的版本，本文就来讨论如何正确设置随机种子。

什么是随机种子

随机数，分为真随机数和伪随机数，真随机数需要自然界中真实的随机物理现象才能产生，而对于计算机来说生成这种随机数是很难办到的。而伪随机数是通过一个初始化的值，来计算来产生一个随机序列，如果初始值是不变的，那么多次从该种子产生的随机序列也是相同的。这个初始值一般就称为种子。

Linux 系统中的随机数

在 Ubuntu 系统中，有一个专门管理随机种子的服务 systemd-random-seed.service，该服务负责在计算机启动的时候，从硬盘上加载一个随机种子文件到内核中，以作为随机初始化值在整个系统运行的过程中提供服务。Linux 会通过许多硬件信息来获得这个初始化值。可以通过 /dev/urandom 文件来产生随机字节，然后使用 od 命令（该命令可将字节转换成希望的格式并打印）来获得随机数：

# 随机生成 0 ~ 255 的数。
# -N1    从 /dev/urandom 中读取一个字节
# -t u1  将字节流格式化为为 1 字节的无符号整型
od -An -N1 -t u1 < /dev/urandom
# 输出 194

# 再次调用
od -An -N1 -t u1 < /dev/urandom
# 输出 50

如果仅希望获得随机数，直接读取 /dev/urandom 或调用 Linux 系统调用 getrandom()(内部也使用 /dev/urandom)是不错的选择。但这种随机数是无法复现的，因为种子是由系统设置的，并且每次开机设置的种子都不一样。在「可复现」的场景中，我们需要的是一种能手动控制随机种子和读取随机序列的方式，以便可以重复获得相同随机序列的功能。

如果一个过程依赖系统产生的随机数，则称这个过程是 Non Deterministic（不确定的）；相反的如果一个过程对相同的输入种子都有相同的输出，则这个随机过程是Deterministic 的。在「可复现」场景中，我们需要保证所有的随机过程都是 Deterministic 的。

/dev/random 可生成「随机性」更强的随机数，但由于其依赖的系统资源更多，导致性能缓慢，因此绝大多数场景都只使用 /dev/urandom。

程序中的随机数

在 PyTorch 中，设置随机种子的方法是 torch.manual_seed(777)，这里 777 就是我们所设置的随机种子，设置完毕后，如果多次调用同样的具有随机过程 PyTorch 方法，就会获得相同的结果，例如下面的代码在多次调用后的打印是一样的：

import torch
torch.manual_seed(777)
print(torch.rand(1))  # 始终输出：tensor([0.0819])
print(torch.rand(1))  # 始终输出：tensor([0.4911])

不论在任何机器或系统，只要使用 torch==1.10.0 版本（其他版本大概率也是 OK 的），输出应该都是长这样的。诶？既然随机种子产生跟系统硬件信息相关，那不同的机器至少应该不一样才对呀？上文说了，在要求「可复现」的场景下，是不能使用 /dev/urandom 来产生随机数的，那剩下的是需要搞清楚 PyTorch 是如何生成随机数的。

通过 torch.manual_seed 方法往下找，可以知道 PyTorch 生成随机数是使用了MT19937（梅森旋轉）算法，这个算法的输入只有一个初始化值也不需要其他的环境信息。因此无论在任何机器，只要 PyTorch 的版本一致（算法部分没有改变）并且设置了随机种子，那么调用随机过程所产生的随机数就是一致的。C++ 11 在标准库中直接引入了这个方法：std::mt19937，而 PyTorch 是自己实现的，官方称性能比 C++ 的版本要更好一些，感兴趣的话可以直接看 PyTorch 源码。

NumPy 的 np.random.seed 也同样使用 MT19937 来生成随机数，因此也与硬件无关。要注意的是：np.random.seed 只影响 NumPy 的随机过程，torch.manual_seed 也只影响 PyTorch 的随机过程。通过下面的代码很容易验证这个结果：

import torch
import numpy

torch.manual_seed(777)

print(torch.rand(1))  # 始终输出 tensor([0.0819])
print(numpy.random.rand(1)) # 多次调用产生不同输出

import torch
import numpy

numpy.random.seed(777)

print(torch.rand(1))  # 多次调用都产生不同输出
print(numpy.random.rand(1)) # 始终输出 [0.15266373]

由此可以得到这样的结论：程序中所有依赖 MT19937 算法产生随机数的包，都需要手动设置随机种子，才能使整个程序的随机性是可复现的。

根据文档，设置 torch.manual_seed 是对所有的设备设置随机种子。目前似乎没有单独为 CPU 设备设置随机种子的方法。

CUDA 的随机数

PyTorch 中，还有另一个设置随机种子的方法：torch.cuda.manual_seed_all，从名字可知这是设置显卡的随机种子。

在 PyTorch 的内部，使用 CUDA Runtime API 提供的 curand 来设置随机种子，根据 curand 的文档，他们提供的所有随机数生成算法都是 Deterministic 的。

import torch
torch.cuda.manual_seed_all(777)

print(torch.rand(1))  # 多次调用都产生不同输出
print(torch.rand(1, device="cuda:0"))  # 始终输出 tensor([0.3530], device='cuda:0')
print(torch.rand(1, device="cuda:1"))  # 始终输出 tensor([0.3530], device='cuda:0')

上面的代码看起来不够「随机」，因为在不同的 GPU 设备上产生了相同的结果，如果希望不同设备可以产生不同的随机数，可以这么做：

import torch

seed = 777
torch.manual_seed(seed)
for i in range(1, torch.cuda.device_count() + 1):
  torch.cuda.set_device(i)
  torch.cuda.manual_seed_all(seed + i)

print(torch.rand(1))  # 始终输出 tensor([0.0819])
print(torch.rand(1, device="cuda:0"))  # 始终输出 tensor([0.4315], device='cuda:0')
print(torch.rand(1, device="cuda:1"))  # 始终输出 tensor([0.6701], device='cuda:1')

上面的代码既保证了随机性（不同设备产生不同的随机数），也保证了确定性（多次调用只产生相同结果）。在真实场景中，一般只会用相同的设备来产生随机数，因此torch.manual_seed(777)应该就能满足大多数需求。

不同设备之间的随机数

先问一个问题：「用 GPU 训练的实验结果，可以在 CPU 上复现吗？」。

答案是「也许可以」。

根据前文可知，CPU 设置随机种子是用 PyTorch 官方实现的 MT19937，而 GPU 是用到了 CUDA Runtime API 的 curand。因此两套实现是完全不同的，那么对于相同的随机种子，理应产生不同的随机序列，用下面的代码可以验证：

import torch

torch.manual_seed(777)
print(torch.rand(1))  # 输出 tensor([0.0819])

import torch

torch.manual_seed(777)
# 将下面的 cuda:0 改为 cuda:1 会产生相同的结果，因为都是 curand 算法
print(torch.rand(1, device="cuda:0"))  # 输出 tensor([0.3530], device='cuda:0')

从上面的例子中知道，对于同一个随机种子，在 CPU 和 GPU 上产出的结果是不同的，因此这种情况在 GPU 上的结果是无法在 CPU 上复现的。那为什么的答案是「也许可以」呢？

因为很多代码，都会在 CPU 上创建 Tensor，再切换到 GPU 上。只要不直接在 GPU 上创建随机变量，就可以避免这个问题。请看下面的例子：

import torch

torch.manual_seed(777)
print(torch.rand(1).to("cuda:0"))  # 输出 tensor([0.0819], device='cuda:0')

上面的代码输出值跟 CPU 一致，但是 device 是在 CUDA 上。这样写可能性能不如直接在 GPU 上直接创建随机变量，但为了保证程序的确定性，牺牲一点性能我认为是值得的。

多进程的随机性

PyTorch 的 torch.utils.data.DataLoader 在 num_worker > 0 的情况下会 fork 出子进程，而通常又会在加载数据的时候做很多「随机变换」，那么就有必要讨论一下多进程下的随机性是怎样的，

子进程一般会保留父进程的一些状态，这也包括随机种子。因此若不做特殊处理，所有子进程都会产生和父进程相同的随机序列。请看下面的例子：

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

np.random.seed(777)

class Random(Dataset):

    def __getitem__(self, index):
        return torch.from_numpy(np.random.rand(4))

    def __len__(self):
        return 4

loader = iter(DataLoader(
    Random(),
    num_workers=2,
    batch_size=2,
))

loader_result = torch.cat([
    next(loader),
    next(loader),
])

print(loader_result)

# 输出
# tensor([[0.1527, 0.3024, 0.0620, 0.4599],
#         [0.8353, 0.9270, 0.7270, 0.7685],
#         [0.1527, 0.3024, 0.0620, 0.4599],
#         [0.8353, 0.9270, 0.7270, 0.7685]], dtype=torch.float64)

可以发现两次 batch 输出的结果是一样的，这是因为主进程中 numpy 的随机性，被两个 worker 保留了，因此两个 worker 的随机性是相同的。

上面的结果在 torch>=1.9.0 是不能复现的，因为 PyTorch 1.9 之后给 DataLoader 默认给每个 worker 重新设置随机种子。

因此我们需要为每一个 worker 设置不同的随机种子以保证随机性，但每次运行又必须要设置相同的随机种子来保证确定性，更好的代码实现如下：

...

seed = 777

def set_seed(seed):
    torch.manual_seed(seed)
    np.random.rand(seed)

set_seed(seed)

loader = iter(DataLoader(
    Random(),
    num_workers=2,
    batch_size=2,
    worker_init_fn=lambda k: set_seed(seed + k + 1)
))

...

这段代码将主进程的随机种子设置为 777，两个 worker 分别设置为 778 和 779。因为每次运行随机种子的值是一样的，因此可以保证确定性，另外每一个 worker 包括主进程的随机种子都不一样，因此随机性也保证了。

类似的，对于分布式训练，也需要做类似的操作，这里考虑单机多卡的情况：

...

seed = 777
rank = torch.distributed.get_rank()
rank_seed = seed + (rank * 1000)

set_seed(rank_seed)

loader = iter(DataLoader(
    Random(),
    num_workers=2,
    batch_size=2,
    worker_init_fn=lambda k: set_seed(rank_seed + k + 1)
))

...

假设有两个 GPU 进行训练，那么第一个 GPU 的主进程和两个 worker 进程的 seed 为：777，778，779；第二块 GPU 是：1777，1778，1779。

Seed Everything

最后点题，祭出一个 seed_everything：

import torch
import numpy as np
import random

seed = 777

def seed_everything(seed):
    if seed >= 10000:
        raise ValueError("seed number should be less than 10000")
    if torch.distributed.is_initialized():
        rank = torch.distributed.get_rank()
    else:
        rank = 0
    seed = (rank * 100000) + seed

    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)

seed_everything(seed)

DataLoader(
  ...
  worker_init_fn=lambda k: seed_everything(seed + (k * 10000))
)