PyTorch : TypeError: flatten() takes at most 1 argument (2 given)

#1

I am trying to run this program in PyTorch which is custom:

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten(start_dim=1)
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(2142, 51),
            nn.ReLU(),
            nn.Linear(51, 1)
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().to(device)

The above code is the custom NeuralNetwork. I defined the training loop below.

def train_loop(data, y, model, loss_fn, optimizer):
    for i in range(data.shape[0]):
        # Compute model prediction and loss
        pred = model(data[i, :, :])
        loss = loss_fn(pred, y[i, :])

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
        print("Loss: {}".format(loss.item()))

The following is how I would like to train it

final_data = torch.randn(500, 42, 51)
final_output = torch.randn(500, 1)
learning_rate = 1e-3
batch_size = 1
epochs = 5
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(final_data, final_output, model, loss_fn, optimizer)

print('Done!!')

The variable final_data is of shape 500x42x51. The variable final_output is of shape 500x1.
I have been trying to run the above data for 10 epochs but I always end up with this error:

Epoch 1
-------------------------------
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-17-3fb934698ecf> in <module>()
      9 for t in range(epochs):
     10     print(f"Epoch {t+1}\n-------------------------------")
---> 11     train_loop(final_data, final_output, model, loss_fn, optimizer)
     12 
     13 print('Done!!')

4 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/flatten.py in forward(self, input)
     38 
     39     def forward(self, input: Tensor) -> Tensor:
---> 40         return input.flatten(self.start_dim, self.end_dim)
     41 
     42     def extra_repr(self) -> str:

TypeError: flatten() takes at most 1 argument (2 given)

The output is basically a classification between 0 or 1.
I am still a newbie in terms of PyTorch and would like some help solving this issue and understanding what’s wrong.

Thank you

#2

Your current code runs into a shape mismatch:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (42x51 and 2142x51)

After fixing it via:

        self.linear_relu_stack = nn.Sequential(
            nn.Linear(51, 51),vvvvv
            ...

you will run into another shape mismatch in the loss calculation:

ValueError: Expected input batch_size (42) to match target batch_size (1).

which is expected, as you are using a fixed batch size of 42 for the model inputs.

Could you describe the data shapes a bit more and how they should be used, please?

#3

Thank you for your time.

The input is basically a 500x42x51. I want the Neural Network to basically take every 42x51 and classify that input as 0 or 1(output). Does it make sense?

Here, I have randomly created a 500x42x51 matrix using torch. I have the actual input data of the same shape.

The output is of shape 500x1 consisting of either 0’s or 1’s. For each 42x51 input there is an output 0 or 1.

#4

Since my data was originally a numpy.ndarray, I changed the data using torch.from_numpy(). I realized I was able to solve the error by this. But it gives me a new error:

RuntimeError: expected scalar type Float but found Double

#5

PyTorch modules expect inputs with the batch dimension in dim0 (there are some exceptions in e.g. RNNs, but this is not interesting for your model). Since you are indexing the [500, 42, 51] tensor in dim0 and are passing the input as [42, 51] the batch size would be 42. Is this expected?

from_numpy() keeps the dtype and since numpy uses float64 by default, you would need to transform it to float32 (the default in PyTorch) via tensor = tensor.float() or transform your model parameters via model.double() (I would use the former approach, as float64 is usually not needed in ML/DL and would slow down your code).