薛定谔的准确率：PyTorch随机数引发的可复现性陷阱

©作者 | 清川

单位 | 上海交通大学博士生

研究方向 | 联邦学习与端云协同推断

本文主要讨论 PyTorch 模型训练中的两种可复现性：一种是在完全不改动代码的情况下重复运行，获得相同的准确率曲线；另一种是改动有限的代码，改动部分不影响训练过程的前提下，获得相同的曲线。

第一种情况，浅显地讲，我们只需要固定所有随机数种子就行

我们知道，计算机一般会使用混合线性同余法来生成伪随机数序列。在我们每次调用 rand() 函数时，就会执行一次或若干次下面的递推公式：

当、和满足一定条件时，可以近似地认为序列中的每一项符合均匀分布，通过我们可以得到 0 到 1 之间的随机数。这类算法都有一个特点，就是一旦固定了序列的初始值，整个随机数序列也就固定了，这个初始值就被我们称作种子。也就是说，我们在程序的起始位置设定好随机数种子，程序单次执行中第次调用到 rand() 得到的数值将会是固定的，一旦程序中 rand() 函数的调用顺序固定，无论程序重复运行多少遍，结果都将是稳定的。在 Minecraft 中我们可以通过特定的种子生成一模一样的世界就是这个原理。

在深度学习中，我们常用 Dropout 减轻过拟合现象，在训练时会随机抑制一定比例的神经元（将激活值设定为零）；常用 RandomFlip、RandomCrop 等方法处理训练集，引入一些随机噪声来提高模型泛化能力；常用 shuffle 的方式从训练集中随机抽取 batch，一方面可以稳定训练，一方面也可以减轻过拟合。这些方法都引入了训练的随机性。我们在炼丹调参的时候肯定希望特定的超参数对应固定的性能，否则就不能肯定模型效果是超参数带来的还是随机性带来的了。

在 PyTorch 中我们一般使用如下方法固定随机数种子。这个函数的调用尽量放在所有 import 之后，其他代码之前。

defseed_everything(seed):
    torch.manual_seed(seed)       
# Current CPU
    torch.cuda.manual_seed(seed)  
# Current GPU
    np.random.seed(seed)          
# Numpy module
    random.seed(seed)             
# Python random module
    torch.backends.cudnn.benchmark = 
False# Close optimization
    torch.backends.cudnn.deterministic = 
True# Close optimization
    torch.cuda.manual_seed_all(seed) 
# All GPU (Optional)

有些工具库中已经给出了类似的函数，但效果需要自己实验确定，比如 pytorch_lightning.seed_everything 中就没有去除 cudnn 对于卷积操作的优化，很多情况下仍然无法复现。建议使用上面给出的代码，至少在我的实验中一直是可以实现稳定复现的。

第二种情况，总的来说，一定要万分确定改动的代码没有影响random()的调用顺序

重复运行的可复现性早有讨论，但修改代码的可复现性其实是更大的陷阱。如果你觉得，这么简单的问题会有人犯错吗？连自己的代码有没有影响训练都不知道吗？我们看如下问题：在固定随机数种子的前提下，你写了一个训练模型的代码，输出了训练的 loss 和准确率并绘制了图像。突然你想在每轮训练之后再测一下测试准确率，于是小心翼翼地修改了代码，那么问题来了，训练的 loss 和准确率会和之前一样吗？

如果你没有加入额外的操作，答案是一定会不一样！我最近的实验中就发现，模型测试的次数会很明显地影响准确率本身，测的次数不一样，准确率也不一样，有时候训练结束的效果甚至会波动 1% 这么大。我实验中要验证的算法是两个模型协同训练的，其中一个模型应该与 Baseline 性能曲线完全相同，现在实验结果却差了 1%，尴尬了！

▲ 海森堡测不准原理

首先排除其他因素，比如我们在测试时确定使用了 model.eval()，避免了前向传播时 Dropout 层起作用，也避免了 BatchNorm 层对数据的均值方差进行滑动平均，可以认为我们避免了一切直接影响模型参数的操作。那究竟是什么在作祟？

首先要清楚我提到的固定随机数种子对可复现性起作用的前提：rand() 函数调用的次序固定。也就是说，假如在某次 rand() 调用之前我们插入了其他的 rand() 操作，那这次的结果必然不同。

>>> 
import torch

>>> 
from utils 
import seed_everything


>>> seed_everything(
0)

>>> torch.rand(
5)

tensor([
0.4963, 
0.7682, 
0.0885, 
0.1320, 
0.3074])


>>> seed_everything(
0)

>>> _ = torch.rand(
1)

>>> torch.rand(
5)

tensor([
0.7682, 
0.0885, 
0.1320, 
0.3074, 
0.6341])

我们再反思一下，模型测试中唯一不敢确定的就是 DataLoader 了。按照常规设置，训练时一般使用带 shuffle 的 DataLoader，而测试时使用不带 shuffle 的，那既然不带 shuffle，为啥还是会出错？我们写一个最小样例复现一下这个问题：

import torch

from torch.utils.data 
import TensorDataset, DataLoader

from utils 
import seed_everything


seed_everything(
0)

dataset = TensorDataset(torch.rand((
10, 
3)), torch.rand(
10))

dataloader = DataLoader(dataset, shuffle=
False, batch_size=
2)

print(torch.rand(
5))

# tensor([0.5263, 0.2437, 0.5846, 0.0332, 0.1387])

seed_everything(
0)

dataset = TensorDataset(torch.rand((
10, 
3)), torch.rand(
10))

dataloader = DataLoader(dataset, shuffle=
False, batch_size=
2)

for inputs, labels 
in dataloader:

pass
print(torch.rand(
5))

tensor([
0.5846, 
0.0332, 
0.1387, 
0.2422, 
0.8155])

然后研读一下 Pytorch 中 DataLoader 的源码就会发现问题所在。

Python 的 in 操作符会先调用后面的迭代器中的 __ iter __ 魔法函数。每次遍历数据集时，DataLoader 的 __ iter __ () 都会返回一个新的生成器，无论上次遍历是否中途 break，它都会重新从头开始。这个生成器底层有一个 _index_sampler，shuffle 设置为真时它使用 BatchSampler（RandomSampler），随机抽取 batchsize 个数据索引，如果为假则使用 BatchSampler（SequentialSampler）顺序抽取。

上面所说的生成器的基类叫做 _BaseDataLoaderIter，在它的初始化函数中唯一调用了一次随机数函数，用以确定全局随机数种子。

class_BaseDataLoaderIter(object):
def__init__(self, loader: DataLoader) -> None:
        ...

        self._base_seed = torch.empty((), dtype=torch.int64).random_(generator=loader.generator).item()

        ...

这里的 _base_seed 将会是一个长整型标量随机数。这个种子会在哪里使用呢？目前只在其子类 _MultiProcessingDataLoaderIter 中使用。当我们将 DataLoader 的 worker 数量设置为大于 0 时，将使用多进程的方式加载数据。在这个子类的初始化函数中会新建 n 个进程，然后将 _base_seed 作为进程参数传入：

...

w = multiprocessing_context.Process(

    target=_utils.worker._worker_loop,

    args=(self._dataset_kind, self._dataset, index_queue,

          self._worker_result_queue, self._workers_done_event,

          self._auto_collation, self._collate_fn, self._drop_last,

          self._base_seed, self._worker_init_fn, i, self._num_workers,

          self._persistent_workers))

w.daemon = 
True
w.start()

...

worker 进程内部实际使用到这个种子的地方如下

def_worker_loop(dataset_kind, dataset, index_queue, data_queue, done_event,

                 auto_collation, collate_fn, drop_last, base_seed, init_fn, worker_id,

                 num_workers, persistent_workers):
    ...

    seed = base_seed + worker_id

    random.seed(seed)

    torch.manual_seed(seed)

if HAS_NUMPY:

        np_seed = _generate_state(base_seed, worker_id)

import numpy 
as np

        np.random.seed(np_seed)

    ...

这些操作将会在 init_fn 之前，控制每个进程起始的随机数种子。但据我观察这些操作已经在 RandomSampler 初始化之后了，所以不知道它们是怎么解决serendipity：可能 95% 的人还在犯的 PyTorch 错误这篇文章提到的低版本 PyTorch 中 DataLoader 随机序列重复的问题的。但这些不是重点，按照 PyTorch 向后兼容的设计理念，这里无论谁继承 _BaseDataLoaderIter 这个基类，无论子类是否用到 _base_seed 这个种子，随机数函数都是会被调用的。调用关系梳理如下：

for inputs, labels 
in DataLoader(...):

pass
# in操作符会调用如下
DataLoader()

    DataLoader.self.__iter__()

        DataLoader.self._get_iterator()

            _MultiProcessingDataLoaderIter(DataLoader.self)

                _BaseDataLoaderIter(DataLoader.self)

                    _BaseDataLoaderIter.self._base_seed = torch.empty(

                        (), dtype=torch.int64).random_(generator=DataLoader.generator).item()

# 一般来说generator是None，我们不指定，random_没有from和to时，会取数据类型最大范围，这里相当于随机生成一个大整数

那么如何解决呢？我尝试过使用 DataLoader 的 generator 参数去指定一个随机数序列，但发现这样只会屏蔽遍历数据操作以外的随机数调用的影响。也就是说，这种情况下，只要调用 DataLoader 的次数变化，还是无法复现。那么最简单有效的方法就是在每次 DataLoader 的 in 操作调用之前都固定一下随机数种子。

defstable(dataloader, seed):
    seed_everything(seed)

return dataloader


for inputs, labels 
in stable(DataLoader(...), seed):

pass

这里需要格外注意的是，stable 函数会使训练时每个 epoch 内部的 shuffle 规律相同！之前我们提到 shuffle 训练集可以减轻模型过拟合，是至关重要的，当每个 epoch 内部第 i 个 batch 的内容都对应相同时，模型会训不起来。所以，一个简单的技巧，在传入随机数种子的时候加上一个 epoch 序号。

for epoch 
in range(MAX_EPOCH):  
# training
for inputs, labels 
in stable(DataLoader(...), seed + epoch):

pass

这时随机数种子的设定和 in 操作绑定成了类似的原子操作，所有涉及到 random() 调用的新增代码都不会影响到准确率曲线的复现了。

本文未讨论的其他随机性

按照本文所说的方法就一定能实现可复现性了吗？不一定。因为随机性还体现在方方面面：比如超参数，当我们改变 DataLoader 的 worker 数量时，显然会引入随机性；比如系统配置，同样的代码在不同架构和精度的 CPU、GPU 上运行，底层优化或者截断误差都可能带来随机性。

在我之前做硬件工作的时候，电池电量不同都可能导致同一个程序在同一块板子上跑出完全不同的结果。可复现性其实是学术界广泛关注的一个专门的研究领域，本文只是为日常模型训练提供一些直观的技巧。

2022年7月1日更新

评论区有位朋友指出 PyTorch 官方给出了一种消除随机性的方法，参考 Reproducibility-DataLoader，以及后续的改进封装 SeedDataLoader，但这份代码并不能解决本文提到的第二种随机性。我先来说说这份代码本来是用来解决什么问题的。

根据官方文档给出的说明，PyTorch 的 Data Loader采用了 Randomness in multi-process data loading 的 reseed 算法，即在每次调用 data loader 实例的 __iter__() 方法时，会通过多进程迭代器创建 n 个 worker，每个 worker 的随机种子由主进程传入 base_seed，然后通过 base_seed+worker_id 的方式生成子进程专属的种子。

上面这些文档应该没有及时更新，只提到这里设置了 torch 自己的随机数，但从本文前面给出的源码中可以看到，它也包含了对 random 库和 numpy.random 库的随机种子的初始化。但假如我们定制化地改写了 DataLoader，用到与这些库独立的其他随机算法库，且我们的程序使用 fork 方法创建子进程时，如果不加设置，这些子进程的随机性就会完全相同。

因此，PyTorch 设置了 worker_init_fn 这个函数，我们可以定制化地在这个函数中设定额外的随机函数库的种子。将这个函数作为 DataLoader 的参数传入，然后它会在前面所说的 torch 自带的子进程随机种子初始化的后面执行，也就是说，除了补充额外的种子，我们也可以选择覆盖掉 torch 的设定。

这个问题在本文提到的 serendipity：可能 95% 的人还在犯的 PyTorch 错误这篇文章中已经详细讲解，在 1.9 以上的 PyTorch 版本中已经修复。由于刚刚提到 torch 已经在子进程中做过常用随机库的初始化了，所以使用文档中给出的例程是多此一举。可能是本文用来验证问题的实验不清晰，这里再给一段验证代码，来说明以上文档中的例程不能解决我们的问题。

import random, numpy, torch

from torch.utils.data 
import DataLoader, TensorDataset

from utils 
import seed_everything


seed_everything(
0)

BATCH_SIZE, NUM_WORKERS = 
8, 
4
dataset = TensorDataset(torch.rand((
100, 
3)), torch.rand(
100))


g = torch.Generator()

g.manual_seed(
0)


defseed_worker(worker_id):
    worker_seed = torch.initial_seed() % 
2 ** 
32
    numpy.random.seed(worker_seed)

    random.seed(worker_seed)


train_data = DataLoader(

    dataset, shuffle=
True,

    batch_size=BATCH_SIZE, num_workers=NUM_WORKERS,

    worker_init_fn=seed_worker, generator=g)

test_data = DataLoader(

    dataset, shuffle=
False,

    batch_size=BATCH_SIZE, num_workers=NUM_WORKERS,

    worker_init_fn=seed_worker, generator=g)


deftest():
for inputs, labels 
in test_data:

pass

deftrain():
for inputs, labels 
in train_data:

        print(labels)

break

if __name__ == 
"__main__":

# case 1
# Result: tensor([0.8174, 0.1753, 0.5049, 0.8947, 0.8472, 0.2588, 0.2568, 0.7127])
    train()


# case 2
# Result: tensor([0.8947, 0.1753, 0.7802, 0.2161, 0.9094, 0.7335, 0.3245, 0.6152])
    test()

    train()