Alright, I put together a small profiler comparison between the different methods we talked about. It’s also totally possible my profiler is not timing things most accurately. Please feel free to correct me if I’m doing something incorrectly.
The profiler compares:
(1) torch.nn.Conv2D on a BP x C x H x W tensor
(2) My custom patch_im2col CUDA kernel on a B x C x P x H x W tensor
(3) torch.nn.unfold + torch.matmul on a B x P x C x H x W tensor
I made a GitHub repo 6 so you can actually run it locally or inspect the code, if you’re interested in trying that.
The forward pass results averaged over 20 iterations on a single NVIDIA GTX 1080T:
Parameters:
B = 1
P = 20480
H = 4
W = 4
in_channels = 32
out_channels = 64
kernel_size = 3
padding = 1
stride = 1
dilation = 1

Various patch convolution methods:
  Using-unfold avg. time: 0.058497 seconds
  Rolled-into-batch avg. time: 0.089629 seconds
  Custom-patch-im2col avg. time: 0.036741 seconds


B = 1
in_channels = 32
H = 600
W = 600
Compare to traditional convolution with B x C x H x W inputs with similar number (actually more) of elements:
  nn.Conv2d: 0.037487 seconds

This test doesn’t really differentiate in an obvious way, so maybe it’s the backward pass that differentiates the speed. I am not 100% sure how to profile that though. Perhaps either of you could shed some light? @albanD @ptrblck