【内部処理】畳み込みの高速処理を実現する方法

畳み込みをforループを用いずにどのように実装されているのかが気になっていたのですが，調べてみると，どうやらnn.Unfoldというメソッドを使うことで行列演算へと帰着させる方法が存在するようです．そこでnn.Unfoldを用いて畳み込みを実装してみました．
（尚，PyTorchが実際にこのような方法を用いてnn.Conv2d等の畳み込み演算を実装しているかまでは確認していません．あくまでそういった実装も可能というだけです）

検証のため，nn.Unfoldを用いた実装による畳み込みの出力と，通常通りnn.Conv2dを介した出力を比較してみました．

実装

batch_size，in_channel，out_channelは任意に変更可能です．

batch_size = 2
in_channel = 3
out_channel = 2

input = torch.FloatTensor(np.random.rand(batch_size, in_channel, 3, 3))
input2 = torch.nn.Unfold(kernel_size=(2, 2), stride=(1, 1), padding=(0, 0))(input)

c = nn.Conv2d(in_channel, out_channel, 2)
print(f'Conv2d層の重み:\n{c.weight.data}\nバイアス:\n{c.bias.data}')

w = c.weight.data.detach().clone()
w = w.view(out_channel, -1)
b = c.bias.data
print('\nカーネルを展開')
print(w)

print('\n展開した入力とカーネルで行列積を行い，整形')
print((torch.matmul(w, input2) + b.view(1, out_channel, 1)).view(batch_size, out_channel, 2, 2))

print('\n通常の畳み込みによる処理')
print(c(input))

出力

Conv2d層の重み:
tensor([[[[ 0.1203,  0.1741],
          [-0.2618,  0.2749]],

         [[ 0.0496,  0.2472],
          [-0.1281,  0.0517]],

         [[ 0.0583,  0.2524],
          [ 0.2381,  0.2859]]],


        [[[ 0.2665, -0.1610],
          [ 0.1499,  0.1978]],

         [[ 0.0339,  0.1021],
          [ 0.0370,  0.0104]],

         [[ 0.1742,  0.2682],
          [-0.0080, -0.2463]]]])
バイアス:
tensor([-0.2232,  0.1847])

カーネルを展開
tensor([[ 0.1203,  0.1741, -0.2618,  0.2749,  0.0496,  0.2472, -0.1281,  0.0517,
          0.0583,  0.2524,  0.2381,  0.2859],
        [ 0.2665, -0.1610,  0.1499,  0.1978,  0.0339,  0.1021,  0.0370,  0.0104,
          0.1742,  0.2682, -0.0080, -0.2463]])

展開した入力とカーネルで行列積を行い，整形
tensor([[[[0.6216, 0.1176],
          [0.6010, 0.4215]],

         [[0.8527, 0.7857],
          [0.6273, 0.5581]]],


        [[[0.7947, 0.3132],
          [0.3201, 0.7330]],

         [[0.6220, 0.8391],
          [0.3936, 0.4128]]]])

通常の畳み込みによる処理
tensor([[[[0.6216, 0.1176],
          [0.6010, 0.4215]],

         [[0.8527, 0.7857],
          [0.6273, 0.5581]]],


        [[[0.7947, 0.3132],
          [0.3201, 0.7330]],

         [[0.6220, 0.8391],
          [0.3936, 0.4128]]]], grad_fn=<ThnnConv2DBackward>)

2つの出力が等しいことから，畳み込みはnn.Unfoldを用いて実装できることがわかりました．

おまけ（nn.Linear）

線形変換のnn.Linearもtorch.matmulで超簡単に実装できます．

in_features = 4
out_features = 2
dim = 3
l = nn.Linear(in_features, out_features)
input = torch.randn(1, dim, in_features)
l(input), torch.matmul(input, l.weight.data.T) + l.bias.data

出力

(tensor([[[-0.6326,  0.3161],
          [-0.0485,  0.2758],
          [-0.4873,  0.5317]]], grad_fn=<AddBackward0>),
 tensor([[[-0.6326,  0.3161],
          [-0.0485,  0.2758],
          [-0.4873,  0.5317]]]))