`Softmax` before (or) after Loss calculation

palguna_gopireddy: <p>I am a basic question.<br>
Should <code>softmax</code> be applied after or before Loss calculation. I have seen many threads discussing the same topic about <code>Softmax</code> and <code>CrossEntropy Loss</code>. But my question is in general, i.e. regarding using <code>Softmax</code> with any loss function. So</p>
<ol>
<li>Is it a rule of thumb that softmax if used, it should only be used before ( or after) loss calculation.</li>
<li>If it is not a rule of thumb, which gives better results. Applying before or After Loss calculation.</li>
</ol>
<p>Like in the below code, Should the <code>Softmax</code> be applied at <strong>Line 1</strong> (or) <strong>Line 2</strong></p>
<pre><code class="lang-auto">_softmax = torch.nn.Softmax(dim = 1)
for epoch in range(num_epochs):
    for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode
            num_data = 0
            num_corrects = 0
            _loss = []
            for i, data in enumerate(tqdm(dataloader['phase'])):
                # Every data instance is an input + label pair
                inputs, true_labels = data

                # Zero your gradients for every batch!
                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    # Make predictions for each batch
                    predictions = model(inputs)

    Line 1: ------- predictions = _softmax(predictions) ---------------       
                    # Compute the loss for each batch
                    loss = loss_fn(predictions, labels)
    Line 2: ------- predictions = _softmax(predictions) ----------------

                    # Calculate predicted labels for each batch
                    _, pred_labels = torch.max(predictions.data, 1)
            
                    if phase == 'train':
                        #Compute loss gradients
                        loss.backward()
                        # Adjust learning weights
                        optimizer.step()

                _loss.append(loss.item())

                # Calculate samples for each epoch
                num_data += true_labels.size(0)

                # Calculate number of correctly predicted labels for each epoch
                num_corrects += torch.sum(pred_labels == true_labels.data.to(device))
            epoch_loss = numpy.mean(_loss)
            epoch_accuracy = 100 * num_corrects / num_data
</code></pre>
<p>(Tell me if there is any step wrong while calculating <code>epoch_loss</code> or <code>epoch_accuracy</code>)</p>

ptrblck: <ol>
<li>It depends on the loss function and where it defines the softmax operation. The issue with <code>F.softmax</code> and <code>nn.CrossEntropyLoss</code> in PyTorch is that <code>nn.CrossEntropyLoss</code> applies <code>F.log_softmax</code> internally and thus no previous <code>F.softmax</code> operation should be used. It’s thus not the user’s choice if and where to use the softmax, but the loss function definition.</li>
<li>Same as before: applying the unnecessary <code>F.softmax</code> before the internal <code>F.log_softmax</code> of <code>nn.CrossEntropyLoss</code> is wrong. Some users reported beneficial results which could have been also achieved by e.g. lowering the learning rate.</li>
</ol>

palguna_gopireddy: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>. You are the best. Clear cut answer.</p>

ptrblck: <p>Sure, happy to help!<br>
One addition: you are of course totally free to still apply <code>F.softmax</code> on the model’s output in case you want to use the probabilities in some debugging step to visualize these. You should however not pass these probabilities to e.g. <code>nn.CrossEntropyLoss</code>.</p>