设置 torch.backends.cudnn.benchmark=True 将会让程序在开始时花费一点额外时间,为整个网络的每个卷积层搜索最适合它的卷积实现算法,进而实现网络的加速。适用场景是网络结构固定(不是动态变化的),网络的输入形状(包括 batch size,图片大小,输入的通道)是不变的,其实也就是一般情况下都比较适用。反之,如果卷积层的设置一直变化,将会导致程序不停地做优化,反而会耗费更多的时间。
添加的位置:
在模型的开始之前
import torch.backends.cudnn as cudnn
cudnn.deterministic = True
cudnn.benchmark = True
使用的测试模型是 ResNet-101,GPU 是 GTX 1060,训练时间缩短了15%