需要pytorch帮助一次完成N个核的N个图像的2D卷积吗

我有一个火炬张量，它是一堆图像。让我们说它是踢

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

这是一个

4x5x6

张量，表示垂直堆叠的四个

5x6

图像

我想用它自己的2-D内核来卷积每一层，这样

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

显然，我可以用for循环来实现这一点，但我想利用GPU的加速，同时完成所有的卷积运算。无论我尝试什么，我只能使用torch的

conv2d

或

conv3d

生成一个输出层，它是所有2d卷积的总和。或者我可以制作4层，每个层都是所有2d卷积的总和。这里有一个具体的例子。让我们使用上面定义的

im

。假设内核由

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.]]])

从上面看，

im

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.]]])

如果我做for循环，正确的答案很简单：

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.]]]])

正如我前面提到的，我唯一能得到的是这四个输出图像的一个总和（或同一个总和层的四个副本）：

---

我确信我想做的是可能的，我只是还没能把我的头绕过去。有人提供解决方案吗？

如果使用分组卷积，这是非常简单的

从

在

groups=in_channels

时，每个输入通道都与自己的通道进行卷积 一套过滤器

这正是我们想要的

需要考虑

F.conv2d

的

weights

参数的形状，因为它根据

组的值而变化。
权重
的第一个维度应该是
输出通道
，在本例中为4。根据的第二个维度应为\u通道/组中的
，即1。因此，我们可以使用

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

它产生一个具有值的形状张量
[4,3,4]

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.]]])

非常感谢你。这也让分组变得非常清楚！
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.]]])