How to apply linear layer to 2D layer only in one dimension (by row or by column) - partially connected layers

Question

I'm trying to apply a linear layer to a 2D matrix of tensors connecting it only by column as in the picture below.

The input shape is (batch_size, 3, 50). I first tried with 2D convolution, adding a 1 channel dimension, so input shape is (batch_size, 1, 3, 50)

import torch.nn as nn
import torch

class ColumnConv(nn.Module):
    def __init__(self):
        self.layers = nn.Sequential(
            nn.Conv2d(
                in_channels=1,
                out_channels=1,
                kernel_size=(3, 1),
                stride=1,
            ),  # shape is B, 1, 1, 50
            nn.ReLU(),
            nn.Flatten() #shape is B, 50
        )

    def forward(self, x):
        return self.layers(x)

But it doesn't seem to work. I'm planning to use a list of 50 nn.Linear layers and apply them to column slices of input, but it seems much more like a workaround not optimized for performance.

Is there a more "pytorchic" way of doing this?

Answer 1

The PyTorch nn.Linear module can be applied to multidimensional input, the linear will be applied to the last dimension so to apply by column the solution is to swap rows and columns.

linear_3_to_1 = nn.Linear(3, 1)

x = torch.randn(1, 1, 3, 50)
x = x.transpose(2, 3) #swap 3 and 50
out = linear_3_to_1(x).flatten()