背景

画像処理でフィルタを変えて同一画像にそれぞれ異なったフィルタ N 個を適用(conv2d)したい
scipy.signal.convolve だと遅い(mode=fft で早くすることできるが, 近似的な結果になりそう...?)ので, pytorch で処理する(CUDA が使えれば GPU で高速化が見込める)
モノクロ画像とする

depthwise convolution(Conv2d で groups を設定) が必要かと思いましたが, in_channel=1, out_channel=N で OK でした.

手順

weight を [N, 1, KH, KW] で作り用意します.

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')


kernel_size = [3, 3]

num_kernels = 4

kernel = torch.zeros([num_kernels, 1, kernel_size[1], kernel_size[0]], dtype=torch.float32).to(device)

# テスト用に channel ごとに all 1s, 2s, 3s, ... なのを用意する
for i in range(num_kernels):
    kernel[i] = float(i) + 1

print(kernel)

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


        [[[2., 2., 2.],
          [2., 2., 2.],
          [2., 2., 2.]]],


        [[[3., 3., 3.],
          [3., 3., 3.],
          [3., 3., 3.]]],


        [[[4., 4., 4.],
          [4., 4., 4.],
          [4., 4., 4.]]]], device='cuda:0')

あとは以下のようにして conv2d を設定 + 実行します!


# assume default stride and dilation and odd kernel size 
conv = torch.nn.Conv2d(1, num_kernels, kernel_size, padding=[kernel_size[1]//2, kernel_size[0] // 2], bias=False, padding_mode='replicate')
conv.weight = torch.nn.Parameter(kernel)

img = np.ones(img_size).astype('float32')

x = torch.from_numpy(img).to(device)

# [h, w] -> [1, 1(c), h, w]
x = torch.unsqueeze(torch.unsqueeze(x, 0), 0)

# forward pass(eval convolution)
with torch.no_grad():
    f = conv(x)

f = f.to('cpu')
print(f)

tensor([[[[ 9.,  9.,  9.,  ...,  9.,  9.,  9.],
          [ 9.,  9.,  9.,  ...,  9.,  9.,  9.],
          [ 9.,  9.,  9.,  ...,  9.,  9.,  9.],
          ...,
          [ 9.,  9.,  9.,  ...,  9.,  9.,  9.],
          [ 9.,  9.,  9.,  ...,  9.,  9.,  9.],
          [ 9.,  9.,  9.,  ...,  9.,  9.,  9.]],

         [[18., 18., 18.,  ..., 18., 18., 18.],
          [18., 18., 18.,  ..., 18., 18., 18.],
          [18., 18., 18.,  ..., 18., 18., 18.],
          ...,
          [18., 18., 18.,  ..., 18., 18., 18.],
          [18., 18., 18.,  ..., 18., 18., 18.],
          [18., 18., 18.,  ..., 18., 18., 18.]],

         [[27., 27., 27.,  ..., 27., 27., 27.],
          [27., 27., 27.,  ..., 27., 27., 27.],
          [27., 27., 27.,  ..., 27., 27., 27.],
          ...,
          [27., 27., 27.,  ..., 27., 27., 27.],
          [27., 27., 27.,  ..., 27., 27., 27.],
          [27., 27., 27.,  ..., 27., 27., 27.]],

         [[36., 36., 36.,  ..., 36., 36., 36.],
          [36., 36., 36.,  ..., 36., 36., 36.],
          [36., 36., 36.,  ..., 36., 36., 36.],
          ...,
          [36., 36., 36.,  ..., 36., 36., 36.],
          [36., 36., 36.,  ..., 36., 36., 36.],
          [36., 36., 36.,  ..., 36., 36., 36.]]]])

Voila!

Channel(layer)ごとに異なった kernel が適用されています!

TODO

RGB 画像で試す