PyTorch中的torch.backends.cudnn.benchmark的作用

最新推荐文章于 2024-08-15 16:43:28 发布

敲代码的小风

最新推荐文章于 2024-08-15 16:43:28 发布

阅读量2.2k

点赞数 3

分类专栏：零基础学习SSD网络PyTorch实现 Deep-Learning-with-PyTorch 《深度学习之PyTorch实战计算机视觉》文章标签：深度学习神经网络 pytorch

本文链接：https://blog.csdn.net/m0_46653437/article/details/110731647

版权

零基础学习SSD网络PyTorch实现同时被 3 个专栏收录

293 篇文章 26 订阅

订阅专栏

Deep-Learning-with-PyTorch

216 篇文章 19 订阅

订阅专栏

《深度学习之PyTorch实战计算机视觉》

198 篇文章 41 订阅

订阅专栏

参考链接: 巧用PyTorch中的torch.backends.cudnn.benchmark减少训练时间

总结:
设置 torch.backends.cudnn.benchmark=True 将会让程序在开始时花费一点额外时间，为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速。适用场景是网络结构固定（不是动态变化的），网络的输入形状（包括 batch size，图片大小，输入的通道）是不变的，其实也就是一般情况下都比较适用。反之，如果卷积层的设置一直变化，将会导致程序不停地做优化，反而会耗费更多的时间。

一般代码放在开头:

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