最靠谱的pytorch训练可复现性设置方案

Leo606

已于 2023-05-16 19:52:51 修改

阅读量1.3k

点赞数 4

文章标签： pytorch 深度学习训练可复现训练

于 2023-05-16 19:48:29 首次发布

本文链接：https://blog.csdn.net/Brouce__Lee/article/details/130712524

版权

原文地址：最靠谱的pytorch训练可复现性设置方案

作者：LeoLi6 (转载引用请注明出处)

保证训练结果的一致和可复现，在算法迭代训练过程中非常重要，否则由于随机性非常容易导致判断错误和干扰算法优化方向。

本文适用于单机多卡的情况。

如果有帮助到你，欢迎点赞收藏加关注！

一、Pytorch seed setting method

1、随机种子

2、训练使用不确定的算法

3、数据加载DataLoader

4、其他特殊情况：lstm dropout

5、Pytorch可复现性完整设置

二、Pytorch-lightning seed setting method

一、Pytorch seed setting method

版本信息：

torch: 1.8.2/1.11.0

影响可复现的因素主要有这几个：

1、随机种子

固定的随机种子是保证可复现性最常用的手段，其中包括random、numpy、以及PyTorch自身的随机种子等，如基本种子、cuda种子、多gpu种子等，此外还需要固定环境变量中的PYTHONHASHSEED。

# seed init.
random.seed(seed)
np.random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)

# torch seed init.
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

2、训练使用不确定的算法

使用的不确定算法算法主要包括两部分：

2.1 CUDA卷积优化——CUDA convolution benchmarking

torch.backends.cudnn.benchmark = False

当torch.backends.cudnn.benchmark选项为True时候，cuda为了提升训练效率，会自动试运行不同优化的卷积算法，以搜索最优最快的算法实现，由于不同硬件不同以及不同的版本的卷积算法实现，可能会导致训练结果不一致。所以，为了算法可复现，通常设置cudnn.benchmark = False。

那什么情况可以设置True:

不考虑可复现性，当模型的输入和结构在训练过程保持固定不变化的时候，可以实现算法加速。

否则，会因为反复的算法最优搜索导致额外的时间浪费。

2.2 Pytorch使用不确定算法——Avoiding nondeterministic algorithms

虽然禁用CUDA卷积基准测优化可以确保CUDA每次运行应用程序时选择相同的卷积算法，但其他算法本身可能是不确定的，如gather等操作，所以需要设置成固定的：

torch.backends.cudnn.deterministic = True

# https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'

# avoiding nondeterministic algorithms (see https://pytorch.org/docs/stable/notes/randomness.html)
torch.use_deterministic_algorithms(True)

torch.use_deterministic_algorithms(True)允许你配置PyTorch在可用的情况下使用确定性算法而不是非确定性算法，并且如果已知某个操作是不确定的(并且没有确定的替代方法)，则抛出RuntimeError错误。

请查看torch.use_deterministic_algorithms()的文档，以获取受影响操作的完整列表。如果操作没有按照文档正确运行，或者如果您需要确定实现没有文档的操作，也可提交问题:Issues · pytorch/pytorch

例如，运行torch.Tensor.index_add_()的不确定CUDA实现将抛出错误:

>>> 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.index_select 这个接口，可以有如下的实现。

def deterministic_index_select(input_tensor, dim, indices):
    """
    input_tensor: Tensor
    dim: dim 
    indices: 1D tensor
    """
    tensor_transpose = torch.transpose(x, 0, dim)
    return tensor_transpose[indices].transpose(dim, 0)

此外还有一种解决办法，可以尝试升级torch版本到更高版本如2.0.1。

3、数据加载DataLoader

在多进程数据加载算法中，DataLoader将根据Randomness in multi-process data loading algorithm对worker进行重新初始化种子。使用worker_init_fn()来保持可再现性:

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

DataLoader(
    xxx,
    batch_size=batch_size,
    num_workers=num_workers,
    worker_init_fn=seed_worker,
)

4、其他特殊情况：lstm dropout

在某些版本的CUDA中，RNN和LSTM网络可能具有不确定性行为，如lstm中的dropout，需要注意这一特性。可以参见torch.nn.RNN()和torch.nn.LSTM()了解详细信息和解决方案。

5、Pytorch可复现性完整设置

最终，一个完整的可复现性设置如下：

def set_seed(seed):
    # seed init.
    random.seed(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)

    # torch seed init.
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.enabled = False # train speed is slower after enabling this opts.

    # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'

    # avoiding nondeterministic algorithms (see https://pytorch.org/docs/stable/notes/randomness.html)
    torch.use_deterministic_algorithms(True)

set_seed(11)

值得注意的是，

如果没有使能torch.backends.cudnn.enabled = False 这一行，我这边测试发现是无法保证训练结果可复现性。但添加这行后会导致训练速度很慢。很多博客没有讲到或者解决这一问题。
此外，如果CUDA版本是10.2或更高版本，少数CUDA操作是不确定的，除非设置了环境变量CUBLAS_WORKSPACE_CONFIG=:4096:8或CUBLAS_WORKSPACE_CONFIG=:16:8。了解更多细节https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility。如果没有设置这些环境变量配置之一，当使用CUDA张量调用这些操作时将引发RuntimeError:

方法优缺点：

第一种方法保证了可复现但会较大降低训练速度。我测试了很多办法，终于解决了这一问题，见第二部分。

二、Pytorch-lightning seed setting method

pytorch-lightning是对pytorch的封装，省去很多繁琐的中间过程，使用起来非常方便。为了确保可复现性，pl提供了设置接口，非常简洁方便。

版本情况：

pytorch-lightning                  2.0.2
torch                              2.0.1

设置：

from lightning.pytorch import Trainer, seed_everything

seed_everything(42, workers=True)
# sets seeds for numpy, torch and python.random.

trainer = Trainer(deterministic=True)

To ensure full reproducibility from run to run you need to set seeds for pseudo-random generators, and set deterministic flag in Trainer.
By setting workers=True in seed_everything(), Lightning derives unique seeds across all dataloader workers and processes for torch, numpy and stdlib random number generators. When turned on, it ensures that e.g. data augmentations are not repeated across workers.

设置workers=True会确保在dataloader中的多线程的可复现性, 并且确保各线程得到各自独特的不会重复的种子，避免因为相同的种子导致相同的数据增强。

deterministic决定pytorch是否使用确定性算法，默认False, 设置True会使pytorch使用确定性算法。

通过上述pl设置，主要有以下优点（重点！）：

保证了可复现的同时，没有牺牲任何训练速度
设置非常简单方便

PyTorch Lightning Reproducibility 2.0.2 documentation

References:

Alxander：PyTorch可复现/重复实验的相关设置
 torch.backends.cudnn.enabled - PyTorch 2.0 documentation
xiaopl：torch.backends.cudnn.benchmark ?!
有关于pytorch模型训练的可复现性_set_deterministic_Reza.的博客-CSDN博客