PyTorch 18. torch.backends.cudnn

写在开头

torch.backends.cudnn.benchmark=True

当该参数设为True时，网络会在初始时寻找最适合当前网络操作的卷积计算方法，能够提高网络的训练效率；但是当输入图像尺寸不断发生变化时，采用该参数会减慢网络训练速度。

准则

如果网络的输入数据维度或类型上变化不大（也就是对输入数据初始化时同一数据尺寸），设置该参数能够增加运行效率；

如果网络的输入数据在每次iteration都变化的话，会导致cuDNN每次寻找一遍最优配置，这样反而会降低运行效率。

背景知识

cuDNN

cuDNN是英伟达专门为深度神经网络所开发的GPU加速库，针对卷积、池化等常见操作做了非常多的低层优化，比一般的GPU程序要快很多。大多数主流深度学习框架都支持cuDNN。在使用GPU的时候，PyTorch会默认使用cuDNN加速。但是，在使用cuDNN的时候，torch.backends.cudnn.benchmark模式为False。

卷积运算

卷积层是卷积神经网络中最重要的部分，也是占用运算量最大的部分，如果可以在低层代码中提升卷积运算效率的话，就可以在不改变给定的神经网络架构的情况下，提升网络训练的速度。

卷积的实现方法有：

1. 直接法，循环多层嵌套
2. GEMM（General Matrix Multiply）
3. FFT(快速傅里叶变换)，先转换到频域，进行相乘再转换到时域
4. 基于Winograd算法
5. …

torch.backends.cudnn.benchmark

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

放置位置

if args.use_gpu and torch.cuda.is_available():
    device = torch.device('cuda')
    torch.backends.cudnn.benchmark = True
else:
    device = torch.device('cpu')