Нужна помощь Pytorch для создания 2D сверток из N изображений с N ядрами одновременно

Question

У меня есть тензор факела, который представляет собой стопку изображений. Скажем, для ударов это

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

, который является 4x5x6 тензором, что означает четыре 5x6 изображения, сложенные вертикально.

Я хочу свернуть каждый слой с его собственным двумерным ядром, чтобы

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

Я, очевидно, могу сделать это с помощью a для l oop, но я бы хотел взять Преимущество ускорения GPU и делать все свертки одновременно. Независимо от того, что я пытаюсь, я смог использовать только факел conv2d или conv3d для получения одного выходного слоя, который является суммой всех 2-х сверток. Или я могу сделать 4 слоя, каждый из которых представляет собой одинаковую сумму всех 2-х сверток. Вот конкретный пример. Давайте использовать 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.]]])

Правильный ответ прост, если я сделаю для l oop:

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

Как я отмечал ранее, я смог получить только одну сумму из этих четырех выходных изображений (или четырех копий одного и того же суммированного слоя):

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

Я уверен, что то, что я хочу сделать, возможно, я просто еще не смог обернуть это вокруг себя. Кто-нибудь может предложить решение?

Answer 1

Это довольно просто, если вы используете групповую свертку.

Из документации nn.Conv2d

В groups=in_channels, каждый входной канал свернут с собственным набором фильтров

Это именно то, что мы хотим.

Необходимо учитывать форму аргумента weights для F.conv2d, так как он изменяется в зависимости от на значение groups. Первое измерение weights должно быть просто out_channels, что в данном случае равно 4. Второе измерение в соответствии с F.conv2d docs должно быть in_channels / groups, что равно 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.]]])