PyTorch Linear layer input dimension mismatch

user_3179768 0
Viewed 6137

Im getting this error when passing the input data to the Linear (Fully Connected Layer) in PyTorch:

matrices expected, got 4D, 2D tensors

I fully understand the problem since the input data has a shape (N,C,H,W) (from a Convolutional+MaxPool layer) where:

  • N: Data Samples
  • C: Channels of the data
  • H,W: Height and Width

Nevertheless I was expecting PyTorch to do the "reshaping" of the data form:

  • [ N , D1,...Dn] --> [ N, D] where D = D1*D2*....Dn

I try to reshape the Variable.data, but I've read that this approach is not recommended since the gradients will conserve the previous shape, and that in general you should not mutate a Variable.data shape.

I am pretty sure there is a simple solution that goes along with the framework, but i haven't find it.

Is there a good solution for this?

PD: The Fully connected layer has as input size the value C * H * W

pytorch deep-learning machine-learning python
3 Answers

A more general solution (would work regardless of how many dimensions x has) is to take the product of all dimension sizes but the first one (the "batch size"):

n_features = np.prod(x.size()[1:])
x = x.view(-1, n_features)
user_102043540

It is common to save the batch size and infer the other dimension in a flatten:

batch_size = x.shape[0]
...
x = x.view(batch_size, -1)
user_90676150
Related