Issue with nn.Softmax() not applying to defined model ouput

peepeepoopoo: <p>Hi! I recently ran into this issue where when using torch.nn.Softmax() under model init, I would find my output classifier without softmax applied.</p>
<p>So in this case, when defining model <strong>init</strong>:</p>
<pre><code class="lang-auto">class TestNet(nn.Module):
  def __init__(self, num_classes=1000 ):
          super(TestNet,self).__init__()
          ...
          ## fully conected layer with output size defined by number of classes
          self.fc = nn.Sequential(
              nn.Linear(256, num_classes),
              nn.Softmax(dim=1)
          )

  def forward(self,x):
       ...
      x = self.fc(x)
      return x
</code></pre>
<p>The result of the output would seem to look as if softmax isn’t applied whatsoever:</p>
<p>But if I apply torch.nn.functional.softmax(output), the result returned is as expected, with values between 0-1, and the sum of all values in each batch equating to 1.</p>
<blockquote>
<p>PS: Note my Criterion is set to use CrossEntropyLoss, which by default should be using torch.nn.LogSoftmax(). Would then in this case effect the nn.Softmax layer applied to the model in init?</p>
</blockquote>

KFrank: <p>Hi Xinan!</p>
<aside class="quote no-group" data-username="peepeepoopoo" data-post="1" data-topic="203207" data-full="true">
<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/peepeepoopoo/48/28687_2.png" class="avatar"> peepeepoopoo:</div>
<blockquote>
<pre><code class="lang-auto">class TestNet(nn.Module):
  def __init__(self, num_classes=1000 ):
          super(TestNet,self).__init__()
          ...
          ## fully conected layer with output size defined by number of classes
          self.fc = nn.Sequential(
              nn.Linear(256, num_classes),
              nn.Softmax(dim=1)
          )

  def forward(self,x):
       ...
      x = self.fc(x)
      return x
</code></pre>
<p>The result of the output would seem to look as if softmax isn’t applied whatsoever:</p>
</blockquote>
</aside>
<p>Works for me.  Here is a script that tests your <code>TestNet</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

import torch.nn as nn

class TestNet(nn.Module):
    def __init__(self, num_classes=1000 ):
        super(TestNet,self).__init__()
        # ...
        ## fully conected layer with output size defined by number of classes
        self.fc = nn.Sequential(
            nn.Linear(256, num_classes),
            nn.Softmax(dim=1)
        )
    
    def forward(self,x):
        # ...
        x = self.fc(x)
        return x

tn = TestNet (num_classes = 4)

print ('tn = ...')
print (tn)

inp = torch.randn (1, 256)

print ('tn (inp) = ...')
print (tn (inp))

print ('tn.fc[0] (inp) = ...')
print (tn.fc[0] (inp))
print ('tn.fc[0] (inp).softmax (dim = 1) = ...')
print (tn.fc[0] (inp).softmax (dim = 1))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.3.0
tn = ...
TestNet(
  (fc): Sequential(
    (0): Linear(in_features=256, out_features=4, bias=True)
    (1): Softmax(dim=1)
  )
)
tn (inp) = ...
tensor([[0.4014, 0.2170, 0.1677, 0.2139]], grad_fn=&lt;SoftmaxBackward0&gt;)
tn.fc[0] (inp) = ...
tensor([[ 0.7785,  0.1635, -0.0946,  0.1489]], grad_fn=&lt;AddmmBackward0&gt;)
tn.fc[0] (inp).softmax (dim = 1) = ...
tensor([[0.4014, 0.2170, 0.1677, 0.2139]], grad_fn=&lt;SoftmaxBackward0&gt;)
</code></pre>
<aside class="quote no-group">
<blockquote>
<blockquote>
<p>PS: Note my Criterion is set to use CrossEntropyLoss, which by default should be using torch.nn.LogSoftmax().</p>
</blockquote>
</blockquote>
</aside>
<p>Yes, <code>CrossEntropyLoss</code> does have <code>LogSoftmax</code> built into it, so you<br>
do <em>not</em> want a <code>Softmax</code> layer in your model.  Presumably you just want<br>
<code>self.fc = nn.Linear(256, num_classes)</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>Would then in this case effect the nn.Softmax layer applied to the model in init?</p>
</blockquote>
</aside>
<p>No.  <code>TestNet</code> has no idea what happens after it is called.  So it doesn’t know<br>
that its output is being fed into <code>CrossEntropyLoss</code> and the fact that it is won’t<br>
affect <code>TestNet</code>’s <code>Softmax</code> layer.</p>
<p>Best.</p>
<p>K. Frank</p>

peepeepoopoo: <p>Hi K.Frank!</p>
<p>Thanks for confirming that <code>CrossEntropyLoss</code> does apply <code>LogSoftmax</code> during loss calculation! (which also make sense that my model is converging properly still).</p>
<p>It’s interesting that in your case, <code>nn.Softmax</code> did apply to your model output.</p>
<blockquote>
<aside class="quote no-group">
<blockquote>
<p>Would then in this case effect the nn.Softmax layer applied to the model in init?</p>
</blockquote>
</aside>
<p>No. <code>TestNet</code> has no idea what happens after it is called. So it doesn’t know<br>
that its output is being fed into <code>CrossEntropyLoss</code> and the fact that it is won’t<br>
affect <code>TestNet</code>’s <code>Softmax</code> layer.</p>
</blockquote>
<p>That is great to know! although I’ll double check that it’s not due to my Criterion setting. I’ll get back to this post once it’s verified that <code>CrossEntropyLoss</code> is interfering with the <code>nn.Softmax</code>  layer.  If not, then it will be a great mystery waiting to be solved haha.</p>
<p>Thanks again for the response!</p>
<p>Xinan</p>

peepeepoopoo: <p>Update:<br>
So I’ve run the same model again but on CPU, and everything behaves as expected with  <code>nn.Softmax</code></p>
<p>Which then prompt me to double-check my code, and realize that I didn’t push my model to GPU with <code>model.to(device)</code><br>
(I’ve set device to GPU)<br>
After doing so, I’m getting the expected result. The weird thing is the model was able to even run in the first place as input and output are pushed to GPU memory, so they should be <code>tesnor.cuda...</code>, which should cause an object type mismatch when performing operations within the model. hmmm</p>
<p>Xinan</p>