PyTorch使用GPU训练设置随机种子后结果仍不可复现该如何解决？

运行环境 #

• PyTorch版本：1.7.1
• PyTorch Geometric版本：1.6.3
• NVIDIA CUDA版本：11.2
• 开发工具：Jupyter Notebook

已配置的基础种子固定代码 #

device = torch.device('cuda:0')
rand = 123
torch.manual_seed(rand)
torch.cuda.manual_seed(rand)
torch.cuda.manual_seed_all(rand)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(rand)
random.seed(rand)

问题表现 #

• 切换cpu设备训练时，每次运行损失、输出结果完全一致
• 使用cuda:0设备训练时，每次运行结果存在差异

修复步骤 #

步骤1：新增CUDA确定性环境变量配置 #

该配置需要在导入torch、numpy等所有依赖包之前设置，可放在Jupyter Notebook的第一个代码块中：

import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'

该配置用于强制CUDA的cublas等基础算子使用确定性运算逻辑，是PyTorch 1.7+版本CUDA训练可复现的必填配置。

步骤2：启用PyTorch全局确定性算法开关 #

在原有种子固定代码中新增如下配置：

# 若运行报错可替换为 torch.set_deterministic(True)
torch.use_deterministic_algorithms(True)

该配置会强制所有PyTorch内置算子使用确定性实现，同时也会适配PyTorch Geometric 1.6.3版本的大部分图算子的确定性需求。

步骤3：固定数据加载模块的随机性 #

如果使用DataLoader加载训练数据，需要额外配置多进程加载的随机种子，避免数据加载过程引入随机性：

def worker_init_fn(worker_id):
    worker_seed = rand + worker_id
    np.random.seed(worker_seed)
    random.seed(worker_seed)

# 初始化DataLoader时传入对应参数
train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    worker_init_fn=worker_init_fn,
    generator=torch.Generator().manual_seed(rand)
)

特殊情况说明 #

如果完成上述配置后仍存在结果不一致的情况，需要检查代码中是否引入了自定义CUDA扩展、第三方非官方实现的算子，这类算子如果未做确定性适配，会成为随机性来源，可替换为PyTorch官方实现的等效算子规避问题。

内容的提问来源于stack exchange，提问作者MH Lee