title: Need pytorch help doing 2D convolutions of N images with N kernels all at once
body: <p>I have a torch tensor which is a stack of images. Let's say for kicks it is</p>
  
  <pre><code>im=th.arange(4*5*6,dtype=th.float32).view(4,5,6)
  </code></pre>
  
  <p>which is a <code>4x5x6</code> tensor, meaning four <code>5x6</code> images stacked vertically. </p>
  
  <p>I want to convolve each layer with its own 2-D kernel so that </p>
  
  <pre><code>I_{out,j} = k_j*I_{in,j}, j=(1...4)
  </code></pre>
  
  <p>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 <code>conv2d</code> or <code>conv3d</code> 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 <code>im</code> as defined above. Say that the kernel is defined by </p>
  
  <pre><code>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.]]])
  </code></pre>
  
  <p>and from above, <code>im</code> is </p>
  
  <pre><code>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.]]])
  </code></pre>
  
  <p>The right answer is easy if I do the for loop:</p>
  
  <pre><code>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.]]]])
  </code></pre>
  
  <p>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):</p>
  
  <pre><code>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.]]]])
  </code></pre>
  
  <p>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?</p>