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)