pytorch中提高模型可复现性的几个设置

szp    2022-04-19 10:42

最近被模型的可复现问题弄得不堪其扰,现在勉强算是解决了,记录一下

首先是最常规的随机种子设置
python中随机种子
random.seed(seed)
numpy中随机种子
np.random.seed(seed)
torch中随机种子
torch.manual_seed(seed)

部分博客中还记录了给GPU设备设置随机种子的设置,torch.cuda.manual_seed给当前GPU设备设置随机种子,torch.cuda.manual_seed_all给所有GPU设备设置随机种子,但官方文档中已经说明了,torch.manual_seed即可直接对所有设备(包括CPU)设置种子,因此如果不是为了设置不同的随机种子,仅仅为了保证模型可复现性,以下两条无需再进行设置,仅使用torch.manual_seed即可
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
除了随机种子以外,很多博客中常见的设置还有
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark=False关闭了cuda的基准测试,使得卷积运算时选择固定的算法(可能会降低性能,但也可能因为无需多次测试不同算法而提高性能),但要注意的是,官方文档里写的很明白,这并不代表cuda一定会有稳定的结果,因为算法本身就是 nondeterministic 的,因此还需要torch.backens.cudnn.deterministic=True,保证cuda的可复现性。

但仅仅这两条是不够的,这就是很多博客里没有说明的坑,上面两行代码控制的只是cuda的运算,pytorch中本身依然会有很多不确定性操作,具体参见文档说明,想要让pytorch具备可复现性,还需要
torch.use_deterministic_algorithms(True)
这个设置会对pytorch中的所有运算生效,包括cuda,因此如果使用了上面这行代码,可以不用再加torch.backens.cudnn.deterministic=True。如果代码中含有nondeterministic的运算,torch.use_deterministic_algorithms()会抛出RuntimeError,如果不想让它抛出错误,可以加上参数warn_only=True
torch.use_deterministic_algorithms(True, warn_only=True)
官方文档中还提供了一个替代函数torch.set_deterministic_debug_mode(),接收str类型的default、warn、error或int类型的0、1、2参数,分别对应不提示、警告和报错三个级别,但实测中发现这个函数似乎并没有真正设置deterministic的操作,因此最后还是使用了torch.use_deterministic_algorithms()。

除此之外,对于RNN模型,10.1版本的cuda需要设置 CUDA_LAUNCH_BLOCKING=1 ,10.2以上版本的cuda需要设置 CUBLAS_WORKSPACE_CONFIG 环境变量,参考文档说明,设置缓冲区大小为 :16:8 或 :4096:2 (注意前面的冒号)

cuda 10.1
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

cuda 10.2及以上
os.environ['CUBLAS_WORKSPACE_CONFIG']=':16:8'       
os.environ['CUBLAS_WORKSPACE_CONFIG']=':4096:2'

实际测试下来,当buffer大小设置为:4096:2 时,torch本身的deterministic检查会给出警告或错误,当buffer大小设置为:16:8 时,RNN的模型虽然变化幅度不大,但仍然不能稳定的百分百复现,。

对于cuda10.2以上的版本,为了保证deterministic的操作,需要设置 CUBLAS_WORKSPACE_CONFIG 环境变量,参见官方文档英伟达文档,设置缓冲区大小为 :16:8 或 :4096:8 (一样要注意前面的冒号),这里不明白为什么和RNN的两种buffer选择不一致,再结合设置:4096:2 时会报错,可能文档有错误?
os.environ['CUBLAS_WORKSPACE_CONFIG']=':16:8'       
os.environ['CUBLAS_WORKSPACE_CONFIG']=':4096:8'
两者设置其中一个即可,前者可能影响整体性能,后者会略微增加显存占用。经过以上设置后,在自己的不包含rnn的模型中,可以做到稳定复现。
Last Modified: 2022-04-19 10:43
Views: 2.2K

[[total]] comments

Post your comment
  1. [[item.time]]
    [[item.user.username]] [[item.floor]]Floor
  2. Click to load more...
  3. Post your comment