Collapsing consecutive linear layers

Question

I have a neural network with 3 consecutive linear layers (convolution), with no activation functions in between. After training the network and obtaining the weights, I would like to collapse all 3 layers into one layer.

How can this be done in practice, when each layer has different kernel size and stride?

The layers are as follows:

1. Convolution layer with a 3x3 kernel, 5 input channels and 5 output channels (a tensor of size 3x3x5x5), with stride 1 and padding "same"
2. Convolution layer with a 5x5 kernel, 5 input channels and 50 output channels (a tensor of size 5x5x5x50), with stride 2 and padding "same"
3. Convolution layer with a 3x3 kernel, 50 input channels and 50 output channels (a tensor of size 3x3x50x50), with stride 1 and padding "same"

Thanks in advance

Answer 1

Permute the dimensions of the first-layer kernels such that input channels are in the "mini-batch" dimension and output channels are in the "channels" dimension. Apply the second layer to that as if that were an image. Then apply the third layer to the result of that. The final result are kernels of the "collapsed" layer.

Use "full" padding all these operations. If that works roughly correctly (apart from padding), try fixing the padding (probably it should be "same" in the last operation).