Torch.nn.functional.conv2d implementation source code?

I need to override the conv2d of pytorch, I am looking for the source which does the convolution operation to understand how it is done.

@ptrblck can you please help me find the source code? I came across this Convolution.cpp, but I was not successful in finding the exact line where a convolution between parameters and input is performed.

You can find some implementations for CUDA e.g., here: pytorch/ConvolutionMM2d.cu at master · pytorch/pytorch · GitHub although the typical “fast” implementations that are used are dispatched from the cuDNN library and not open-sourced.

On CPU, something similar exists here: pytorch/ConvolutionMM2d.cpp at master · pytorch/pytorch · GitHub although again in practice the “fast” implementations are in MKLDNN/oneDNN.

2 Likes