该系列其他文章

Seed Everything - 可复现的 PyTorch（一）

在上一篇文章中介绍了通过设置随机种子来保证程序的确定性，对于一般的程序来说，这已经足够了，但对于 PyTorch 等深度学习相关的程序，还有其他的随机变量需要固定，特别是跟 CUDA 有关的函数。

cuDNN benchmark

cuDNN 是英伟达开源的深度学习加速软件，其中实现了一些常用操作（例如卷积）。而 PyTorch 的部分方法底层也依赖 cuDNN。

cuDNN 有一个方法 cudnnFindConvolutionForwardAlgorithmEx，即针对某一组固定的卷积参数（包括输入），该方法内部会遍历所有的卷积实现算法（卷积的实现算法不是只有一种），然后根据性能来选择对于该组参数的最佳算法。

这类似于一种自动调优（Auto perf tuning），看上去很美妙，但实际也会有一些缺点：

  1. 如果用户程序中，卷积层的输入或参数经常发生变化，则会频繁触发 Perf tuning，反而会极大降低性能。
  2. 不同显卡设备或硬件可能会让 cudnnFindConvolutionForwardAlgorithmEx 选择出不同的实现算法，而这可能会引入不确定性。

可以用下面这段代码快速验证一下性能的问题，虽然性能不是本文的重点：

import torch
import time

BENCHMARK = True  # or False
DYNAMIC_INPUT = True  # or False

torch.backends.cudnn.benchmark = BENCHMARK

resnet50 = torch.hub.load("pytorch/vision", "resnet50").cuda()
start = time.time()
for i in range(50):
    if DYNAMIC_INPUT:
        inputs = torch.rand(size=(2, 3, *torch.randint(256, 512, size=(2,)).tolist()), device="cuda:0")
    else:
        inputs = torch.rand(size=(2, 3, 512, 512), device="cuda:0")
    resnet50(inputs)

print(f"BENCHMARK={BENCHMARK}, DYNAMIC_INPUT={DYNAMIC_INPUT}, Time Usage: ", time.time() - start)

四次实验的输出：

BENCHMARK=False, DYNAMIC_INPUT=False, Time Usage:  0.6690871715545654
BENCHMARK=True, DYNAMIC_INPUT=False, Time Usage:  0.6754634380340576
BENCHMARK=False, DYNAMIC_INPUT=True, Time Usage:  0.7371523380279541
BENCHMARK=True, DYNAMIC_INPUT=True, Time Usage:  23.754278898239136

可以看见当输入是动态的，并且 benchmark=True 的时候，性能表现非常差；即使输入是固定的，benchmark=True 所带来的性能提升也不明显（测试结果几乎没有提升），更不用说随之而来的不确定性问题。

因此，torch.backends.cudnn.benchmark 设置为 False 在多数情况下应该是更好的选择。

不确定性算法

PyTorch 中有很多不确定的算法，而其根源大多都来自于 CUDA API（cuDNN 和 cuBLAS）。上文提到同一种操作（例如卷积）在 cuDNN 中有多种实现，而不同的实现方式在性能和确定性上都可能不一样。好在我们可以强制 PyTorch 只选择确定性的实现方式。

import torch

torch.use_deterministic_algorithms(True)

至于 cnDNN 部分实现为什么是不确定的，本文不再深入，感兴趣的读者可以参考官方文档。

该代码使 PyTorch 在有多种算法选择的前提下，仅选择确定性的算法；如果某种操作不提供确定性的实现版本，则直接报错。可以在 use_deterministic_algorithms 的文档查到不确定性的方法列表。

如果设置了 torch.use_deterministic_algorithms(True)，需要同时设置环境变量 CUBLAS_WORKSPACE_CONFIG=:4096:8，否则绝大多数模型都会因为不确定性而报错，原因可参考 cuBLAS 文档，这个行为可能会在将来被修复。

总结

如果希望保证程序的确定性，需要

  1. 正确设置随机种子
  2. 使用确定性的算法

把前一篇文章最后抬出的代码完善一下，最终保证确定性的代码如下：

import os
import torch
import numpy as np
import random

# 避免使用不确定性算法
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False  # 可不加，因为缺省值就是 False

# 设置随机种子
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))
)