torch.backends.cudnn.benchmark 和torch.backends.cudnn.deterministic

  • Post author:
  • Post category:其他


torch.backends.cudnn.benchmark标志位True or False

cuDNN是GPU加速库

在使用GPU的时候,PyTorch会默认使用cuDNN加速,但是,在使用 cuDNN 的时候,torch.backends.cudnn.benchmark模式是为False。

设置这个 flag 为True,我们就可以在 PyTorch 中对模型里的卷积层进行预先的优化,也就是在每一个卷积层中测试 cuDNN 提供的所有卷积实现算法,然后选择最快的那个。这样在模型启动的时候,只要额外多花一点点预处理时间,就可以较大幅度地减少训练时间。

如果我们的网络模型一直变的话,不能设置cudnn.benchmark=True。因为寻找最优卷积算法需要花费时间。

这段代码一般放在训练代码的开头,比如再设置使用GPU的同时,加在后面

如果在 PyTorch 程序中设置了 torch.backends.cudnn.deterministic=True,并且 cudnn.benchmark == False的话,那么就选那个默认的卷积算法

torch.backends.cudnn.deterministic 将这个 flag 置为 True 的话,每次返回的卷积算法将是确定的,即默认算法。如果配合上设置 Torch 的随机种子为固定值的话,应该可以保证每次运行网络的时候相同输入的输出是固定的。

cudnnGetConvolutionForwardAlgorithm_v7vscudnnFindConvolutionForwardAlgorithmEx(链接是到英伟达提供的官方 API)。我们已经知道 PyTorch 默认调用的是前者,设置benchmark=True会调用后者。根据官方讲解来说, Get 那个函数会使用一些人为设置的启发式的方法(heuristic)去选择程序所认为的最合适的算法。而 Find 那个函数是穷尽式的 (exhaustive search),即会遍历所有可选的卷积进行比较。这么一说,其实 PyTorch 默认也是会对每层的卷积算法进行预先选择,速度比较快,但是选择出来的结果不是那么好,具体的选择机制并不是很清楚(没找到相关的资料)。

cudnn.benchmark=True的时候,如果不停地改变输入形状,运行效率确实会很低,因为对于每个新遇到的卷积场景 PyTorch 都会去自动寻找最适合的卷积算法。