Can linear layers take multidimensional input?

#1

If each feature I want my neural network to analyse is a 2d array, would this be possible to pass to a linear layer or would it require flattening?

Im fairly new to Neural networks so any help would appreciated

Thanks

#2

Hi,

Yes, it would require flattening if you want to consider all the elements of the array. I have put a small example printing the size of the output when you give a 2D input to a Linear layer without flattening.

    # 2d input
    linear_layer_2d = nn.Linear(in_features=64, out_features=32)
    # 1st dimension (128) = batch dimension, input 64 x 64
    input_2d = torch.randn(128, 64, 64)
    output_2d = linear_layer_2d(input_2d)
    print(output_2d.size())
    # torch.Size([128, 64, 32])
    
    # 1d input (2d flattened)
    linear_layer_1d = nn.Linear(in_features=4096, out_features=32)
    # input_1d size = [128, 4096]
    input_1d = torch.flatten(input_2d, start_dim=1)
    output_1d = linear_layer_1d(input_1d)
    print(output_1d.size())
    # torch.Size([128, 32])

Let me know if it’s not clear!

2 Likes