一、前言
在深度学习实验中，难免引入随机因素，让模型结果产生一定程度的波动，影响模型的可复现性。而可复现性对于研究人员来讲非常重要，它能够让我们控制变量，摒弃随机因素的影响，从而快速有效地验证自己的想法，提高学习和工作效率。本文将介绍在PyToch框架下如何保证模型的可复现性。
- 这里我们仅仅保证模型在同一个环境下（包括不限于Python版，PyTorch版本，cuda版本，显卡型号等，总之就是同一台机器上配置环境不变的情况下），两次运行同一个程序的结果一致，不保证程序移植到其他环境中的可复现性。对于我们来说，前者已经足够满足需求。
- 不同的PyTorch版本设置方式不尽相同，总体趋势是Pytorch版本越高，设置项越多。文章最后会给出包括
1.4.0，1.7.1以及最新的1.9.0在内的三个Pytorch版本对应的建议。由于笔者精力有限，无法给出每一个版本的设置细节，感兴趣的读者可以自己摸索。 - 如果你和我一样，强迫症晚期患者、希望保证模型运行两次的结果在小数点后n位都保持一致、追求极致的对齐，那就赶紧来看一下吧。
二、影响因素
下面我们对影响模型可复现性的几个因素进行介绍。
1 随机种子
保证可复现性最常用的手段就是设置固定的随机种子，包括random、numpy随机种子，以及PyTorch自身的随机种子等，这其中又包括基本种子，cuda种子，多gpu种子等。某些网友还反映需要固定环境变量中的PYTHONHASHSEED。因此，最简单的方法就是一股脑全部设置：
- 随机种子懒人版：
def set_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
np.random.seed(seed) # Numpy module.
random.seed(seed) # Python random module.
os.environ['PYTHONHASHSEED'] = str(seed)
2. DataLoader
既然可复现性是随机因素引起的，那么是不是设置好随机种子就可以保证复现性了呢？当然不是。这个需要分版本讨论，1.4及以前的版本基本可以满足需求，但是1.4之后还需要商榷。其中需要注意的就是Dataloader部分，
在1.7.1版本中
def worker_init(worked_id):
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
train_loader = DataLoader(xxx, num_workers=0, worker_init_fn=worker_init)
具体为啥这样就不说了，见官网 https://pytorch.org/docs/1.8.0/notes/randomness.html?highlight=reproducibility
3. 避免不确定性算法
cuDnn benchmark ,cuDnn中往往会并行跑多个算法来选择效果最好的那个，这里就可能引入噪声。因此，如果为了保证可复现性，可以将这个特性禁用，缺点就是会损失一定的性能。如果为了达到最优的效果，可以将开关打开：
torch.backends.cudnn.benchmark = False # 禁用benchmark，保证可复现
或者：
torch.backends.cudnn.benchmark = True # 恢复benchmark，提升效果
避免原子操作 一些操作使用了原子操作，不是确定性算法，不能保证可复现，因此我们使用下面的代码禁用原子操作，保证使用确定性算法：
torch.set_deterministic(True)
这样的设置可以让模型把某不确定算法转成对应的确定性算法来执行。如果有部分操作没有对应的确定性算法，例如gather, scatter, repeat_interleave等，这样就会报错。
>>> import torch
>>> torch.use_deterministic_algorithms(True)
>>> torch.randn(2, 2).cuda().index_add_(0, torch.tensor([0, 1]), torch.randn(2, 2))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: index_add_cuda_ does not have a deterministic implementation, but you set
'torch.use_deterministic_algorithms(True)'. ..
这个时候，我们有几种选择，
- 放弃使用确定性算法。即取消
torch.set_deterministic(True)，这样就无法模型保证可复现了。 - 手工替换对应的操作。例如，将
gather替换成其他等价的操作，这个还需要具体思考是否有合适的替代方案。 - 向官方报告报告链接，请求添加对应算法的确定性版本，这个不能解燃眉之急。
cuda的确定性，根据英伟达官方的说法，在cuda 10.2及以上的版本中，需要设置以下环境变量来保证cuda的结果可复现(官方链接)：
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
4. 其他
某些模块，如CUDA RNN 和LSTM中可能还有不确定算法，这个需要根据具体情况来决定。
三、完整设置
下面我们根据不同的PyTorch版本给出设置建议：
- Pytorch 1.7.1
种子：
def set_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
np.random.seed(seed) # Numpy module.
random.seed(seed) # Python random module.
torch.set_deterministic(True)
torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
os.environ['PYTHONHASHSEED'] = str(seed)
DataLoader:
def worker_init(worked_id):
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
train_loader = DataLoader(xxx, num_workers=0, worker_init_fn=worker_init)
网友评论