Need pytorch help doing 2D convolutions of N images with N kernels all at once
I have a torch tensor which is a stack of images. Let's say for kicks it is

im=th.arange(4*5*6,dtype=th.float32).view(4,5,6)


which is a 4x5x6 tensor, meaning four 5x6 images stacked vertically. 

I want to convolve each layer with its own 2-D kernel so that 

I_{out,j} = k_j*I_{in,j}, j=(1...4)


I can obviously do this with a for loop, but I'd like to take advantage of GPU acceleration and do all the convolutions at the same time. No matter what I try, I've only been able to use torch's conv2d or conv3d to produce a single output layer that is the sum of all the 2d convolutions. Or I can make 4 layers where each is the same sum of all the 2d convolutions. Here's a concrete example. Let's use im as defined above. Say that the kernel is defined by 

k=th.zeros((4,3,3),dtype=th.float32)
n=-1
for i in range(2):
    for j in range(2):
        n+=1
        k[n,i,j]=1
        k[n,2,2]=1
print(k)

tensor([[[1., 0., 0.],
         [0., 0., 0.],
         [0., 0., 1.]],

        [[0., 1., 0.],
         [0., 0., 0.],
         [0., 0., 1.]],

        [[0., 0., 0.],
         [1., 0., 0.],
         [0., 0., 1.]],

        [[0., 0., 0.],
         [0., 1., 0.],
         [0., 0., 1.]]])


and from above, im is 

tensor([[[  0.,   1.,   2.,   3.,   4.,   5.],
         [  6.,   7.,   8.,   9.,  10.,  11.],
         [ 12.,  13.,  14.,  15.,  16.,  17.],
         [ 18.,  19.,  20.,  21.,  22.,  23.],
         [ 24.,  25.,  26.,  27.,  28.,  29.]],

        [[ 30.,  31.,  32.,  33.,  34.,  35.],
         [ 36.,  37.,  38.,  39.,  40.,  41.],
         [ 42.,  43.,  44.,  45.,  46.,  47.],
         [ 48.,  49.,  50.,  51.,  52.,  53.],
         [ 54.,  55.,  56.,  57.,  58.,  59.]],

        [[ 60.,  61.,  62.,  63.,  64.,  65.],
         [ 66.,  67.,  68.,  69.,  70.,  71.],
         [ 72.,  73.,  74.,  75.,  76.,  77.],
         [ 78.,  79.,  80.,  81.,  82.,  83.],
         [ 84.,  85.,  86.,  87.,  88.,  89.]],

        [[ 90.,  91.,  92.,  93.,  94.,  95.],
         [ 96.,  97.,  98.,  99., 100., 101.],
         [102., 103., 104., 105., 106., 107.],
         [108., 109., 110., 111., 112., 113.],
         [114., 115., 116., 117., 118., 119.]]])


The right answer is easy if I do the for loop:

import torch.functional as F
for i in range(4):
    print(F.conv2d(im[i].expand(1,1,5,6),k[i].expand(1,1,3,3)))

tensor([[[[14., 16., 18., 20.],
          [26., 28., 30., 32.],
          [38., 40., 42., 44.]]]])
tensor([[[[ 75.,  77.,  79.,  81.],
          [ 87.,  89.,  91.,  93.],
          [ 99., 101., 103., 105.]]]])
tensor([[[[140., 142., 144., 146.],
          [152., 154., 156., 158.],
          [164., 166., 168., 170.]]]])
tensor([[[[201., 203., 205., 207.],
          [213., 215., 217., 219.],
          [225., 227., 229., 231.]]]])


As I noted earlier, the only thing I've been able to get is one sum of those four output images (or four copies of the same summed layer):

F.conv2d(im.expand(1,4,5,6),k.expand(1,4,3,3)) 

tensor([[[[430., 438., 446., 454.],
          [478., 486., 494., 502.],
          [526., 534., 542., 550.]]]])


I'm certain that what I want to do is possible, I just haven't been able to wrap my head around it yet. Does anyone have a solution to offer?

This is pretty straight forward if you use a grouped convolution.

From the nn.Conv2d documentation


  At groups=in_channels, each input channel is convolved with its own
  set of filters


Which is exactly what we want.

The shape of the weights argument to F.conv2d needs to be considered since it changes depending on the value of groups. The first dimension of weights should just be out_channels, which is 4 in this case. The second dimension according to F.conv2d docs should be in_channels / groups, which is 1. So we can perform the operation using

F.conv2d(im.unsqueeze(0), k.unsqueeze(1), groups=4).squeeze(0)


which produces a tensor of shape [4,3,4] with values

tensor([[[ 14.,  16.,  18.,  20.],
         [ 26.,  28.,  30.,  32.],
         [ 38.,  40.,  42.,  44.]],

        [[ 75.,  77.,  79.,  81.],
         [ 87.,  89.,  91.,  93.],
         [ 99., 101., 103., 105.]],

        [[140., 142., 144., 146.],
         [152., 154., 156., 158.],
         [164., 166., 168., 170.]],

        [[201., 203., 205., 207.],
         [213., 215., 217., 219.],
         [225., 227., 229., 231.]]])

