【Pytorch使用教程】torch.backends.cudnn.benchmark = True的作用

Source

在 PyTorch 中,设置 torch.backends.cudnn.benchmark = True 是一种优化深度学习应用程序性能的方法,特别是当你有固定输入大小的时候。

解释

CuDNN:CUDA Deep Neural Network library(CuDNN)是 NVIDIA 提供的一个 GPU 加速库,用于深度神经网络。PyTorch 在底层使用 CuDNN 来加速 GPU 上的运算。

cudnn.benchmark:这个设置告诉 CuDNN 自动寻找最适合当前配置的高效算法(对于你的特定网络和硬件架构)。这通常可以增加运行效率,但有两个主要的前提条件:

  1. 固定的输入大小:因为 CuDNN 会为每一种新的输入大小重新寻找最优算法。如果网络的输入数据大小在每次前向传播时都在变(例如,每次输入的图片尺寸不同),设置 benchmark = True 可能会导致性能下降,因为每次变化都需要重新寻找算法。

  2. 额外的启动时间:启动训练时,第一次运行会稍微慢一点,因为 CuDNN 在试图寻找最优的算法。但在这之后,如果输入大小不变,后续的运行会更快。

使用场景

你应该在以下情况考虑启用这个设置:

  • 你正在使用较大的固定大小输入(例如,所有图像都预先调整到相同的尺寸)。
  • 你希望尽可能优化模型的运行时间&#