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

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.

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)

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)

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.