Flatten batch input through linear layer in model

1

hi,
i have input tensor [10,1,74,74] here 10 is batch size,
so i need flatten these vector and get [10,1024] output through fc layer(self.fc = nn.Linear(5476,1024)) like this , how can i do that.(74*74=5476)

input= [10,1,74,74]
out=[10,1024]

2

Hi, if you want to flatten a tensor you could use the .view() function. So in your case

import torch

input_ = torch.ones((10, 1, 74, 74))

print(input_.shape)
# torch.Size([10, 1, 74, 74])

flattened_input = input_.view(10, -1)

print(flattened_input.shape)
# torch.Size([10, 5476])

linear_layer = torch.nn.Linear(5476, 1024).to("cuda")
out = linear_layer(flattened_input.to("cuda"))

print(out.shape)
# torch.Size([10, 1024])