MNIST dataset - Why is my loss so high? [beginner]

username: mgloria
cooked: <p>I wanted to build a simple ANN and train it from scratch on the Mnist dataset. The accuracy values look fine as expected but the loss is just way too high, as if I was not computing it correctly. What am I missing in the code below?</p>
<pre><code class="lang-auto">import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn 
import torch.nn.functional as F
import torch.optim as optim
torch.set_grad_enabled(True) #by default gradients are computed
from itertools import product
from sklearn.metrics import accuracy_score


class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        # hidden layers
        self.h1 = nn.Linear(in_features=784, out_features=196)
        self.h2 = nn.Linear(in_features=196, out_features=49)
        # output layer, 10 units - one for each digit
        self.output = nn.Linear(49, 10)
        
    def forward(self, x):
        # (0) input layer
        x = x
        # (1) first hidden layer
        x = self.h1(x)
        x = F.relu(x)
        # (2) second hidden layer
        x = self.h2(x)
        x = F.relu(x)
        # (3) output layer 
        x = self.output(x)
        # x = F.softmax(x, dim=1) -&gt; already in cross-entropy loss function??
        return x
    
def get_num_correct(preds, labels):
    return preds.argmax(dim=1).eq(labels).sum().item() 

train_set = torchvision.datasets.MNIST(
    root='./data', 
    train=True,
    download=True, 
    transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))

parameters = dict(lr = [.01], batch_size = [32, 64])
param_values = [v for v in parameters.values()]

for lr, batch_size in product(*param_values): 
    comment = f' batch_size={batch_size} lr={lr}'
    print(comment)
    network = Network()
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size)
    optimizer = optim.Adam(network.parameters(), lr=lr)
    images, labels = next(iter(train_loader))
    
    for epoch in range(3):
        total_loss = 0
        total_correct = 0
        all_preds = torch.Tensor([])
        for batch in train_loader:
            images, labels = batch # Get Batch
            preds = network(images.unsqueeze(0).reshape(-1, 28*28*1)) # Pass batch
            all_preds = torch.cat((all_preds, preds), dim=0)
            loss = F.cross_entropy(preds, labels) # Calculate Loss
            optimizer.zero_grad() # Zero Gradients
            loss.backward() # Calculate Gradients
            optimizer.step() # Update Weights

            total_loss += loss.item() * batch_size # is this correct? why? 
            total_correct += get_num_correct(preds, labels)

        print(f"""epoch: {epoch}; total loss: {total_loss}; 
              accuracy: {accuracy_score(train_loader.dataset.targets, all_preds.argmax(dim=1))}""")

</code></pre>

username: ptrblck
cooked: <p>Your <code>total_loss</code> consists of the losses of all samples in your <code>Dataset</code>. Usually you print the average loss per sample. To do this, you could divide <code>total_loss</code> by <code>len(train_set)</code>.</p>

username: mgloria
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="62670">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Usually you print the average loss per sample. To do this, you could divide <code>total_loss</code> by <code>len(train_set)</code> .</p>
</blockquote>
</aside>
<p>Thanks a lot <a class="mention" href="/u/ptrblck">@ptrblck</a>, I never realized about this detail! Can I ask you a couple questions more related questions:</p>
<ol>
<li>
<p>how is softmax activation function already implicit in the cross-entropy loss function?</p>
</li>
<li>
<p>Why in order to compute the total loss I am multiplying loss.item() * batch_size?</p>
</li>
</ol>

username: mailcorahul
cooked: <ol>
<li>
<p><code>nn.CrossEntropy</code> combines <code>log_softmax</code> and <code>NLLLoss</code>, that’s the reason there is no need to use <code>nn.Softmax</code> at the end of the network when <code>nn.CrossEntropy</code> is used as a criterion.<br>
Check this: <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/loss.py#L826" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/loss.py#L826</a></p>
</li>
<li>
<p>Generally average loss over the entire batch is returned after <code>criterion(outputs, targets)</code> call. So inorder to get the sum of all losses over a batch, you do <code>loss.item() * batch_size</code>.</p>
</li>
</ol>