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

uhvardhan93: <p>I am trying to run this program in PyTorch which is custom:</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>The above code is the custom NeuralNetwork. I defined the training loop below.</p>
<pre><code class="lang-auto">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()))
</code></pre>
<p>The following is how I would like to train it</p>
<pre><code class="lang-auto">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!!')
</code></pre>
<p>The variable <code>final_data</code> is of shape 500x42x51. The variable <code>final_output</code> is of shape 500x1.<br>
I have been trying to run the above data for 10 epochs but I always end up with this error:</p>
<pre><code class="lang-auto">Epoch 1
-------------------------------
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
&lt;ipython-input-17-3fb934698ecf&gt; in &lt;module&gt;()
      9 for t in range(epochs):
     10     print(f"Epoch {t+1}\n-------------------------------")
---&gt; 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) -&gt; Tensor:
---&gt; 40         return input.flatten(self.start_dim, self.end_dim)
     41 
     42     def extra_repr(self) -&gt; str:

TypeError: flatten() takes at most 1 argument (2 given)
</code></pre>
<p>The output is basically a classification between 0 or 1.<br>
I am still a newbie in terms of PyTorch and would like some help solving this issue and understanding what’s wrong.</p>
<p>Thank you</p>

ptrblck: <p>Your current code runs into a shape mismatch:</p>
<pre><code class="lang-python">RuntimeError: mat1 and mat2 shapes cannot be multiplied (42x51 and 2142x51)
</code></pre>
<p>After fixing it via:</p>
<pre><code class="lang-python">        self.linear_relu_stack = nn.Sequential(
            nn.Linear(51, 51),vvvvv
            ...
</code></pre>
<p>you will run into another shape mismatch in the loss calculation:</p>
<pre><code class="lang-python">ValueError: Expected input batch_size (42) to match target batch_size (1).
</code></pre>
<p>which is expected, as you are using a fixed batch size of 42 for the model inputs.</p>
<p>Could you describe the data shapes a bit more and how they should be used, please?</p>

uhvardhan93: <p>Thank you for your time.</p>
<p>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?</p>
<p>Here, I have randomly created a 500x42x51 matrix using torch. I have the actual input data of the same shape.</p>
<p>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.</p>

uhvardhan93: <p>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:</p>
<p>RuntimeError: expected scalar type Float but found Double</p>

ptrblck: <p>PyTorch modules expect inputs with the batch dimension in <code>dim0</code> (there are some exceptions in e.g. RNNs, but this is not interesting for your model). Since you are indexing the <code>[500, 42, 51]</code> tensor in <code>dim0</code> and are passing the input as <code>[42, 51]</code> the batch size would be <code>42</code>. Is this expected?</p>
<p><code>from_numpy()</code> keeps the <code>dtype</code> and since numpy uses <code>float64</code> by default, you would need to transform it to <code>float32</code> (the default in PyTorch) via <code>tensor = tensor.float()</code> or transform your model parameters via <code>model.double()</code> (I would use the former approach, as <code>float64</code> is usually not needed in ML/DL and would slow down your code).</p>