How can I implement 2D CNN filter with channelwise-bound kernel weights?

Question

I would like to bind kernel parameters through channels/feature-maps for each filter. In a conv2d operation, each filter consists of HxWxC parameters I would like to have filters that have HxW parameters, but the same (HxWxC) form.

The scenario I have is that I have 4 gray pictures of bulb samples (yielding similar images from each side), which I overlay as channels, but a possible failure that needs to be detected might only appear on one side (a bulb has 4 images and a single classification). The rotation of the object when the picture is taken is arbitrary. Now I solve this by shuffling the channels at training, but it would be more efficient if I could just bind the kernel parameters. Pytorch and Tensorflow solutions are both welcome.

Answer 1

Assuming you want HxWx1 kernel to perform convolution on hxwxc images.

Here's sample code which uses single channel kernel to operate on multichannel feature: maps

import torch
import torch.nn as nn
import torch.nn.functional as F


class model(nn.Module):
    def __init__(self, in_ch=4):
        super().__init__()
        self.in_ch  = in_ch
        # single channel kernel initialization
        self.kernel = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=1, stride=1)

    def forward(self, x):
        bs, ch, w, h = x.shape
        x   = x.view((bs, ch, 1, w, h))
        out = self.kernel(x[:, 0])
        # reusing of same kernel
        for i in range(1, self.in_ch):
            out = torch.cat((out, self.kernel(x[:, i])), 1)
        return out

net = model(4)
print(net)
inp = torch.randn((10, 4, 100, 100))
out = net(inp)
print(out.shape)

(The main hack is in the forward function)

Answer 2

In PyTorch implementation of convolution modules, you can just set the out_channels and groups argument to be the number of your input channels. In your case, if your input has 4 channels, you can simply construct the layer as

torch.nn.Conv2d(in_channels=4, out_channels=4, kernel_size=3, groups=4)

Answer 3

This Keras solution uses a shared sub-model, and then applies it to each of the input image's channels separately. The architecture is very similar to a Siamese network.

from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.layers import Input, InputLayer, Dense, Flatten
from tensorflow.keras.layers import Conv2D, MaxPooling2D

import tensorflow.keras.backend as K

res = 128

sub_model = Sequential([
    InputLayer((res,res,1)),
    Conv2D(16, 5, activation='relu'),
    MaxPooling2D(2),
    Conv2D(16, 5, activation='relu'),
    MaxPooling2D(2),
    Conv2D(16, 5, activation='relu'),
    MaxPooling2D(2),
    Flatten(),
    Dense(512, activation='relu'),
    Dense(128, activation='relu'),
])

sub_model.summary()

Summary:

Model: "sequential_8"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_19 (Conv2D)           (None, 124, 124, 16)      416       
_________________________________________________________________
max_pooling2d_17 (MaxPooling (None, 62, 62, 16)        0         
_________________________________________________________________
conv2d_20 (Conv2D)           (None, 58, 58, 16)        6416      
_________________________________________________________________
max_pooling2d_18 (MaxPooling (None, 29, 29, 16)        0         
_________________________________________________________________
conv2d_21 (Conv2D)           (None, 25, 25, 16)        6416      
_________________________________________________________________
max_pooling2d_19 (MaxPooling (None, 12, 12, 16)        0         
_________________________________________________________________
flatten_7 (Flatten)          (None, 2304)              0         
_________________________________________________________________
dense_4 (Dense)              (None, 512)               1180160   
_________________________________________________________________
dense_5 (Dense)              (None, 128)               65664     
=================================================================
Total params: 1,259,072
Trainable params: 1,259,072
Non-trainable params: 0
_________________________________________________________________

Then the combined model:

n_channels = 4
inp = Input((res,res,n_channels))
x = K.concatenate([sub_model(inp[:,:,:,i:i+1])[:,:,None] for i in range(n_channels)])

model = Model(inp, x)

x.shape
# TensorShape([None, 128, 4])

Naturally it is up to you whether to have those Dense layers in the sub_model or not, and which extra operations you'll do to the concatenated x.