人工智能:Pytorch训练可重复性(Reproducibility)

在不同的 PyTorch 版本或不同平台上,不能保证完全可复现的结果。此外,即使使用相同的随机种子,CPU 和 GPU 执行之间的结果也可能无法复现。但是,您可以采取措施来限制平台、设备和 PyTorch 版本的非确定性行为来源。首先,您可以控制随机性的来源,这些来源会导致应用程序的多次执行表现不同。其次,您可以将 PyTorch 配置为避免对某些操作使用非确定性算法,以便在给定相同输入的情况下,对这些操作的多次调用将产生相同的结果。最后,在执行多进程数据加载时还需注意控制加载的随机性工作进程种子。本文内容参考 Pytorch 官方文档:https://pytorch.org/docs/stable/notes/randomness.html#controlling-sources-of-randomness

Deterministic operations are often slower than nondeterministic operations, so single-run performance may decrease for your model. However, determinism may save time in development by facilitating experimentation, debugging, and regression testing.

控制随机性的来源

PyTorch 随机数生成器

您可以使用 torch.manual_seed() 为所有设备(CPU 和 CUDA)的 RNG 设置种子

1
2
import torch
torch.manual_seed(0)

某些 PyTorch 操作可能会在内部使用随机数。例如,torch.svd_lowrank() 就会这样做。因此,使用相同的输入参数连续多次调用它可能会得到不同的结果。但是,只要在应用程序开始时将 torch.manual_seed() 设置为常量,并且消除了所有其他非确定性来源,则每次在相同环境中运行应用程序时,都会生成相同的随机数序列。

还可以通过在后续调用之间将 torch.manual_seed() 设置为相同的值,从使用随机数的操作中获得相同的结果。

Python 随机数生成器

对于自定义运算符,您可能还需要设置 Python 种子

1
2
import random
random.seed(0)

其他库中的随机数生成器

如果您或您正在使用的任何库依赖于 NumPy,则可以使用以下代码为全局 NumPy RNG 设置种子

1
2
import numpy as np
np.random.seed(0)

但是,某些应用程序和库可能使用 NumPy 随机数生成器对象,而不是全局 RNG(https://numpy.com.cn/doc/stable/reference/random/generator.html),这些对象也需要一致地设置种子。

如果您正在使用任何其他使用随机数生成器的库,请参阅这些库的文档,以了解如何为它们设置一致的种子。

CUDA 卷积基准测试

由 CUDA 卷积操作使用的 cuDNN 库可能是应用程序多次执行之间非确定性的来源。当使用一组新的尺寸参数调用 cuDNN 卷积时,一个可选功能可以运行多个卷积算法,对它们进行基准测试以找到最快的算法。然后,在该过程的剩余时间内,将针对相应的尺寸参数集一致地使用最快的算法。由于基准测试噪声和不同的硬件,即使在同一台机器上,基准测试也可能会在后续运行中选择不同的算法。

使用 torch.backends.cudnn.benchmark = False 禁用基准测试功能会导致 cuDNN 确定性地选择算法,但可能会降低性能。

但是,如果您不需要在应用程序的多次执行之间保持可复现性,则启用基准测试功能 torch.backends.cudnn.benchmark = True 可能会提高性能。

请注意,此设置与下面讨论的 torch.backends.cudnn.deterministic 设置不同。

避免非确定性算法

torch.use_deterministic_algorithms() 允许您将 PyTorch 配置为在可用时使用确定性算法而不是非确定性算法,并在已知操作是非确定性算法(并且没有确定性替代算法)时抛出错误。

有关受影响操作的完整列表,请参阅 torch.use_deterministic_algorithms() 的文档。如果某个操作没有按照文档正确执行,或者您需要一个没有确定性实现的操作的确定性实现,请提交问题:https://github.com/pytorch/pytorch/issues?q=label:%22module:%20determinism%22

例如,运行 torch.Tensor.index_add_() 的非确定性 CUDA 实现将引发错误:

1
2
3
4
5
6
7
>>> 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)'. ...

当使用稀疏密集 CUDA 张量调用 torch.bmm() 时,它通常使用非确定性算法,但当打开确定性标志时,将使用其备用的确定性实现:

1
2
3
4
5
6
7
>>> import torch
>>> torch.use_deterministic_algorithms(True)
>>> torch.bmm(torch.randn(2, 2, 2).to_sparse().cuda(), torch.randn(2, 2, 2).cuda())
tensor([[[ 1.1900, -2.3409],
[ 0.4796, 0.8003]],
[[ 0.1509, 1.8027],
[ 0.0333, -1.1444]]], device='cuda:0')

此外,如果您正在使用 CUDA 张量,并且您的 CUDA 版本为 10.2 或更高版本,则应根据 CUDA 文档设置环境变量 CUBLAS_WORKSPACE_CONFIG:https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility

CUDA RNN 和 LSTM

在某些版本的 CUDA 中,RNN 和 LSTM 网络可能具有非确定性行为。有关详细信息和解决方法,请参阅 torch.nn.RNN()torch.nn.LSTM()

填充未初始化的内存

torch.empty()torch.Tensor.resize_() 等操作可以返回包含未定义值的未初始化内存的张量。如果需要确定性,则使用此类张量作为另一个操作的输入是无效的,因为输出将是非确定性的。但实际上没有任何东西可以阻止运行此类无效代码。因此,为了安全起见,默认情况下 torch.utils.deterministic.fill_uninitialized_memory 设置为 True,如果设置了 torch.use_deterministic_algorithms(True),它将使用已知值填充未初始化的内存。这将防止这种非确定性行为的可能性。

然而,填充未初始化的内存不利于性能。因此,如果您的程序有效并且不使用未初始化的内存作为操作的输入,则可以关闭此设置以获得更好的性能。

DataLoader

DataLoader 将按照 多进程数据加载中的随机性 算法重新设定工作进程的种子。使用 worker_init_fn()generator 来保持可重复性

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2**32
numpy.random.seed(worker_seed)
random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(0)

DataLoader(
train_dataset,
batch_size=batch_size,
num_workers=num_workers,
worker_init_fn=seed_worker,
generator=g,
)

实际实践中的随机种子定义函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import os, gzip, re, string, glob, time, random, shutil, itertools, json

import torch
from torch.cuda.amp import GradScaler, autocast
from torch.nn import CrossEntropyLoss
from torch.optim import SGD, AdamW
from torch.optim.lr_scheduler import LambdaLR,CosineAnnealingLR
from torch.utils.data import IterableDataset, DataLoader, DistributedSampler

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, AutoConfig, set_seed, Trainer

def setup_seed(seed):
"""
Repeatability of model training process. Refer to the following link:
https://pytorch.ac.cn/docs/stable/notes/randomness.html
"""
random.seed(seed)
np.random.seed(seed)

set_seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.use_deterministic_algorithms(True)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"