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

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

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?

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.

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

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).