设置torch.backends.cudnn.benchmark=True后，如何查看及手动指定卷积算法？

开启torch.backends.cudnn.benchmark = True时的算法选择 #

当你开启torch.backends.cudnn.benchmark = True后，cuDNN会自动在FFT、Winograd、GEMM这些卷积算法里跑一遍基准测试，挑出当前硬件和卷积参数（比如输入尺寸、卷积核大小、通道数等）下最快的那个来用。

这几种算法的适用场景大概是：

• GEMM：适合大尺寸卷积核或者大输入张量的情况，计算逻辑直接，但小尺寸场景下效率一般；
• FFT：更适合大输入、小卷积核的场景，靠频域转换来减少计算量；
• Winograd：对中等大小的卷积核（比如3x3）和输入效果好，能减少乘法运算次数，在低算力设备上优势更明显。

cuDNN会根据你的GPU型号、当前卷积层的参数组合自动筛选最优算法，不用你自己判断哪种合适。

手动指定卷积算法的方法 #

PyTorch支持手动指定卷积算法，步骤是先关闭自动基准测试，再给卷积层指定具体算法：

首先关闭自动基准：

import torch
torch.backends.cudnn.benchmark = False

然后给卷积层指定算法，常见的算法枚举在torch.backends.cudnn模块下，比如：

# 创建卷积层
conv = torch.nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3)

# 指定用GEMM算法
conv.algorithm = torch.backends.cudnn.CUDNN_CONVOLUTION_ALGORITHM_GEMM

# 或者指定FFT
# conv.algorithm = torch.backends.cudnn.CUDNN_CONVOLUTION_ALGORITHM_FFT

# 或者指定Winograd
# conv.algorithm = torch.backends.cudnn.CUDNN_CONVOLUTION_ALGORITHM_WINOGRAD

你还可以用torch.backends.cudnn.get_algorithm(conv)来查看当前卷积层用的算法，确认设置是否生效。

需要注意的是，手动指定前要确保你的硬件和cuDNN版本支持该算法，不然可能会报错。

内容的提问来源于stack exchange，提问作者HITLCX