pytorch中_cudnn_convolution_full_forward的定义在哪里

Posted 2021-04-25

正如所提到的标题，我想找到_cudnn_convolution_full_forward的定义，但是我搜索了pytorch中的所有项目并且失败了。我无法找到并讲述这个功能。

任何人都可以帮助我吗？

答案

所有cudnn卷积函数都在这里定义：https://github.com/pytorch/pytorch/blob/1848cad10802db9fa0aa066d9de195958120d863/aten/src/ATen/native/cudnn/Conv.cpp

在最新版本的pytorch中不再存在此功能。最接近的是cudnn_convolution_forward。在版本0.1.12中，该函数位于同一文件中：https://github.com/pytorch/pytorch/blob/v0.1.12/torch/csrc/cudnn/Conv.cpp

我建议不要使用非公共API（一个以_开头）并使用公共方法，但你可能已经知道了。

换句话说，你应该使用

torch.backends.cudnn.enabled = True

然后conv2d或conv3d取决于您的使用。