Compute mse_loss() with softmax()

Mukesh1729: <ul>
<li>Hi I am using using a network that produces an output heatmap (torch.rand(1,16,1,256,256)) with<br>
Softmax( ) as the last network activation.</li>
<li>I want to compute the MSE loss between the output heatmap and a target heatmap.</li>
<li>When I add the softmax the network loss doesn’t decrease and is around the same point and works when I remove the softmax.</li>
</ul>
<p>How can I go about computing mse loss by using softmax()?<br>
Thanks</p>
<pre><code class="lang-auto">from torch.nn import init
class NET(torch.nn.Module):
    def __init__(self):
        super(NET, self).__init__()
        self.encoder = torch.nn.Sequential(
                                          torch.nn.Conv3d(3, 64, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=64),
                                          torch.nn.ReLU(inplace=True),
                                          torch.nn.MaxPool3d(kernel_size=(1,2,2),stride=(1,2,2)),
            
                                          torch.nn.Conv3d(64, 128, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),                                         
                                          torch.nn.BatchNorm3d(num_features=128),
                                          torch.nn.ReLU(inplace=True),
                                          torch.nn.MaxPool3d(kernel_size=(2,2,2),stride=(2,2,2)),
                                          
                                           
                                          torch.nn.Conv3d(128, 256, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=256),
                                          torch.nn.ReLU(inplace=True),
                                          torch.nn.MaxPool3d(kernel_size=(1,2,2),stride=(1,2,2)),
                                         
                                           
                                          torch.nn.Conv3d(256, 512, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=512),
                                          torch.nn.ReLU(inplace=True),
                                          
                                          
                                          )
        
        self.decoder = torch.nn.Sequential(
                                          torch.nn.ConvTranspose3d(512, 256, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1), output_padding=(0,1,1)),
                                          torch.nn.BatchNorm3d(num_features=256), 
                                          torch.nn.ReLU(inplace=True),
                                           
                                                                                  
                                          torch.nn.Conv3d(256, 256, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=256),
                                          torch.nn.ReLU(inplace=True),
                                          
                                          
                                           
                                          torch.nn.ConvTranspose3d(256,128, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1),output_padding=(0,1,1)), 
                                          torch.nn.BatchNorm3d(num_features=128),
                                          torch.nn.ReLU(inplace=True),
                                          
                                           
                                          torch.nn.Conv3d(128, 128, kernel_size=(1,3,3), stride=(1,1,1), padding=(1,1,1)),
                                          torch.nn.BatchNorm3d(num_features=128),
                                          torch.nn.ReLU(inplace=True),
                                          

                                          torch.nn.ConvTranspose3d(128,16, kernel_size=(1,3,3), stride = (1,2,2), padding=(1,1,1),output_padding=(0,1,1)),
                                          )                  

    def forward(self, image):
        
        encoder = self.encoder(image)
        decoder = self.decoder(encoder)
        return torch.nn.Softmax(dim=1)(decoder)
</code></pre>

KFrank: <p>Hi Mukesh!</p>
<aside class="quote no-group" data-username="Mukesh1729" data-post="1" data-topic="137473" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/90db22/48.png" class="avatar"> Mukesh1729:</div>
<blockquote>
<ul>
<li>When I add the softmax the network loss doesn’t decrease and is around the same point and works when I remove the softmax.</li>
</ul>
</blockquote>
</aside>
<p>As a practical matter, I would suggest not using <code>softmax()</code> for the<br>
simple reason that your network trains better without it than with it.</p>
<p>Most of what we know about what works with neural networks comes<br>
not from theory or mathematical proofs, but from experience.  So go<br>
ahead and do what your experience with your use case shows works.</p>
<p>I can give some hand-waving intuition about why <code>softmax()</code> might<br>
not be a good fit with <code>mse_loss()</code>.  The “mse” in <code>mse_loss()</code> stands<br>
for “mean-squared-error.”  Roughly speaking, this is the variance of<br>
the mismatch between your predictions and targets (and the variance<br>
is the square of the standard deviation).  Standard deviations and<br>
variances “naturally” work in a context where the values involved are<br>
unconstrained and run from <code>-inf</code> to <code>inf</code>.  Your <code>softmax()</code> takes<br>
such unbounded values and forces them to be probability-like values<br>
that range over <code>[0.0, 1.0]</code>.</p>
<p>This doesn’t mean that using <code>softmax()</code> couldn’t work or be a good<br>
choice for your use case, but this intuition suggests that <code>softmax()</code><br>
isn’t a natural fit when using <code>mse_loss()</code>.</p>
<p>Best.</p>
<p>K. Frank.</p>

Mukesh1729: <p>Hey Frank, thanks for the tip. I am trying to replicate a paper and they state:</p>
<p><strong>The last convolutional layer of the decoder is followed by a</strong><br>
<strong>transposed-convolutional layer that uses a linear activation followed by a Softmax transformation, resulting in an output</strong><br>
<strong>of 17 confidence maps, one per keypoint</strong>.<br>
ref : <a href="https://arxiv.org/abs/2104.08029" class="inline-onebox" rel="noopener nofollow ugc">[2104.08029] T-LEAP: occlusion-robust pose estimation of walking cows using temporal information</a><br>
When i use softmax as the last layer as suggested in the model then my loss stagnates. It would be nice to know how I can use the softmax (before or after, how should my target heatmap look like then) with the mse loss.  An example would be helpful</p>
<p>For now, i am doing something like this:</p>
<ul>
<li>input = model + torch.nn.Softmax(dim=1) (torch.rand(1,16,2,256,256))</li>
<li>heatmap = torch.rand(1,16,1,256,256)</li>
<li>criterion = torch.nn.MSELoss()</li>
<li>criterion(input, heatmap)</li>
</ul>

KFrank: <p>Hi Mukesh!</p>
<p>I’m not familiar with your use case and I haven’t looked at the paper<br>
you cite, so what I say will be speculation.  Nonetheless:</p>
<aside class="quote no-group quote-modified" data-username="Mukesh1729" data-post="3" data-topic="137473" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/90db22/48.png" class="avatar"> Mukesh1729:</div>
<blockquote>
<p><strong>an output of 17 confidence maps</strong>.</p>
</blockquote>
</aside>
<p>The term “confidence map” (together with the use of <code>Softmax</code>)<br>
suggests to me that you might be dealing with probability-like<br>
values.  If so, cross-entropy (or some other probability-comparison<br>
metric such as the <a href="https://pytorch.org/docs/1.10.0/generated/torch.nn.KLDivLoss.html" rel="noopener nofollow ugc">Kullback-Leibler divergence</a>) might be more<br>
appropriate.</p>
<p>Your discussion suggests that your <code>target</code> heatmap is not made up<br>
of categorical labels, but rather, of continuous probability-like values.</p>
<p>Pytorch’s built-in <code>CrossEntropyLoss</code> does not support such “soft”<br>
labels, although they do make perfect sense.  If you want to explore<br>
the cross-entropy approach, you will have to write your own “soft-label”<br>
version of cross entropy, which is easy enough to do as described here:</p>
<aside class="quote quote-modified" data-post="2" data-topic="69501">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/soft-cross-entropy-loss-tf-has-it-does-pytorch-have-it/69501/2">Soft Cross Entropy Loss (TF has it does Pytorch have it)</a> 
  </div>
  <blockquote>
    Hello Raaj! 

I do not believe that pytorch has a “soft” cross-entropy function built in. 
But you can implement it using pytorch tensor operations, so you should 
get the full benefit of autograd and gpu acceleration. 
See this (pytorch version 0.3.0) script: 
import torch
torch.__version__

# define "soft" cross-entropy with pytorch tensor operations
def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).sum() / input.shap…
  </blockquote>
</aside>

<p>Note, when doing this you still do not want a final <code>Softmax</code> layer.<br>
The <code>input</code> to <code>softXEnt()</code> will be the output of your final <code>Linear</code><br>
layer, understood to be raw-score <em>logits.</em>  The target to <code>softXEnt()</code><br>
will, however, be probability-like values that range over <code>[0.0, 1.0]</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>For now, i am doing something like this:</p>
<ul>
<li>input = model + torch.nn.Softmax(dim=1) (torch.rand(1,16,2,256,256))</li>
<li>heatmap = torch.rand(1,16,1,256,256)</li>
</ul>
</blockquote>
</aside>
<p>I do have some questions about the shapes of your <code>input</code> and<br>
<code>heatmap</code>.</p>
<p>First, what do the dimensions mean?</p>
<p>Second, you mention “17 confidence maps” while the second<br>
dimension of your shapes is <code>16</code>.  Should those values be the same,<br>
or is <code>16</code> an unrelated value with a different meaning?</p>
<p>Last, regardless of the meanings of the dimensions, your <code>input</code> and<br>
<code>heatmap</code> have different shapes, which is logically inappropriate for<br>
<code>MSELoss</code>.  (<code>MSELoss</code> will <em>broadcast,</em> but you probably don’t want<br>
that.)  Why is the third dimension of <code>input</code> <code>2</code>, and is that really what<br>
you want?</p>
<p>Best.</p>
<p>K. Frank</p>

Mukesh1729: <p>Hey Frank,</p>
<ul>
<li>
<p><strong>The term “confidence map” (together with the use of <code>Softmax</code>)</strong><br>
<strong>suggests to me that you might be dealing with probability-like</strong><br>
<strong>values</strong> - Yes,  I am want the probability of a finding a keypoint in a heatmap</p>
</li>
<li>
<p>I mention 16 confidence maps instead of 17 because I am using 16 keypoints - 1 per channel</p>
</li>
<li>
<p>Regarding the input output shape of images they look like below:<br>
Input shape  = torch.rand(1,16,2,256,256) # 2 refers to the depth (in my case images)<br>
model output = torch.rand(1,16,1,256,256)<br>
target shape = torch.rand(1,16,1,256,256)</p>
</li>
</ul>
<p>Considering the last layer of my model is not a softmax layer but a convtranspose layer with no subsequent activation. Does the implementation look right? The values looks large</p>
<pre><code class="lang-auto">def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).sum() / input.shape[0]

input  = torch.rand(1, 16, 1, 256, 256)
target = torch.rand(1, 16, 1, 256, 256)
softXEnt(input, target)

output:tensor(1474121.)
</code></pre>

KFrank: <p>Hi Mukesh!</p>
<p>Three comments:</p>
<aside class="quote no-group" data-username="Mukesh1729" data-post="5" data-topic="137473" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/90db22/48.png" class="avatar"> Mukesh1729:</div>
<blockquote>
<p>Input shape  = torch.rand(1,16,2,256,256) # 2 refers to the depth (in my case images)</p>
</blockquote>
</aside>
<p>First, if this really is the shape of the “image” you pass into your <code>NET</code><br>
model, it won’t work.  From your first post, the initial layer of <code>NET</code> is:</p>
<pre><code class="lang-auto">torch.nn.Conv3d(3, 64, kernel_size=(3,3,3), stride=(1,1,1), padding=(1,1,1))
</code></pre>
<p>That is, your <code>Conv3d</code> layer expects <code>in_channels = 3</code>, and you are<br>
passing in a tensor with <code>in_channels</code> (the second dimension) equal<br>
to <code>16</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).sum() / input.shape[0]

input  = torch.rand(1, 16, 1, 256, 256)
target = torch.rand(1, 16, 1, 256, 256)
softXEnt(input, target)

output:tensor(1474121.)
</code></pre>
</blockquote>
</aside>
<p>Second, as written, the <code>softXEnt()</code> code I posted is taking the mean<br>
over the batch (in your example, batch-size = 1), but is taking the sum<br>
over the image dimensions (in your case, 1x256x256).</p>
<p>(The <code>softXEnt()</code> code I wrote wasn’t directed at the “multidimensional<br>
case” with trailing “image” dimensions.)</p>
<p>If you prefer to average the loss over the image dimensions (which<br>
might be more natural), you could do something like:</p>
<pre data-code-wrap="python"><code class="lang-python">def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).mean()
</code></pre>
<p>(I don’t know why I didn’t just use <code>.mean()</code> in the first place.)</p>
<p>Third, you are using <code>input  = torch.rand(1, 16, 1, 256, 256)</code>.<br>
I realize that this is just an example, but this <code>input</code> consists of values<br>
that range over <code>[0.0, 1.0]</code> (because you are using <code>rand()</code>).  Such<br>
values look like probabilities.  Make sure that your actual <code>input</code> to<br>
<code>softXEnt()</code> (or <code>CrossEntropyLoss</code>, for that matter) consists of <em>logits</em><br>
that run from <code>-inf</code> to <code>inf</code>, and not probabilities.  (If you wanted to make<br>
your example more realistic in this regard, you could use <code>randn()</code> in<br>
place of <code>rand()</code>.)</p>
<p>Best.</p>
<p>K. Frank</p>

Mukesh1729: <p>Hey Frank,</p>
<p>Thanks for getting back. Regarding the dimensions of the image that was a typo, but I have tried what you suggested and it works:</p>
<pre><code class="lang-auto">def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).mean()

# Input image to the model: (1, 3, 2 ,256, 256) 

#Ouput from the model                  
model_out  = torch.randn(1, 16, 1 ,256, 256)
#Target heatmap
target_hmp = torch.randn(1, 16, 1 ,256, 256)

softXEnt(model_out, target_hmp)

output:tensor(-0.0006)
</code></pre>
<p>I added the above loss to my training loop and the loss looks like this. It starts off very low<br>
I am using around 400 training images and 100 validation images<br>
Legends:<br>
<strong>blue: training</strong><br>
<strong>yellow:validation</strong><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/7/9/79da183242e7249fe014b01a01d4bdfb5edb11d3.png" alt="image" data-base62-sha1="hnX2qXA0BJuaheB1Yt0uiJDtnp1" width="263" height="179"></p>
<p><strong>Training loop</strong></p>
<pre><code class="lang-auto">def loss_plot(epochs, train, val):
    ep = [i for i in range(epochs)]
    plt.plot(ep,train,label="Training loss")
    plt.plot(ep,val,label="Validation loss")
    plt.title("Training and validation loss")
    plt.show()
    
def softXEnt (input, target):
    logprobs = torch.nn.functional.log_softmax (input, dim = 1)
    return  -(target * logprobs).mean()



#training and Evaluation
#Get parameters, start training model
DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
EPOCHS     = 10
lr         = 1e-3
model      = T_LEAP()
criterion  = torch.nn.MSELoss()
optimizer  = torch.optim.Adam(model.parameters(),lr=lr, amsgrad=True)

soft       = torch.nn.Softmax(dim=1)
# scheduler = StepLR(optimizer, step_size=5, gamma=0.1)


model = model.to(DEVICE)
tr_losses  = [] 
val_losses = []

tr_batch = int(len(train_loader)/train_loader.batch_size)
vl_batch = int(len(valid_loader)/valid_loader.batch_size)

for idx in range(EPOCHS):
  curr_trloss  = 0.0
  curr_valoss  = 0.0
  model.train()
    
  for i,d in enumerate(train_loader):
      
    # Get images, move them to GPU
    image   = d["images"]
    image   = image.permute(0,2,1,3,4).to(DEVICE)
    heatmap = d["heatmaps"]
    heatmap = heatmap.permute(0,4,1,2,3).to(DEVICE)
    

    optimizer.zero_grad()

    output = model(image)
    loss   = softXEnt(output, heatmap)
  
    loss.backward()
    optimizer.step()
    
    curr_trloss += loss.item()
  tr_losses.append(curr_trloss/tr_batch)
    
 
  # Evaluation loop
  model.eval()
  with torch.no_grad():
      for i,d in enumerate(valid_loader):

        # Get images, move them to GPU
        image   = d["images"]
        image   = image.permute(0,2,1,3,4).to(DEVICE)
        heatmap = d["heatmaps"]
        heatmap = heatmap.permute(0,4,1,2,3).to(DEVICE)

        output = model(image)
        loss   = softXEnt(output, heatmap)

        curr_valoss += loss.item()
        # scheduler.step()
          
  val_losses.append(curr_trloss/vl_batch)
    
  # Print losses 
  print(f"Training loss {idx}:",curr_trloss/tr_batch)
  print(f"Validation loss {idx}:",curr_valoss/vl_batch)
  print(" ")
    
print("Training and evalutation is now complete")
#plotloss    
loss_plot(EPOCHS,tr_losses, val_losses)

#save model
state = {
          'epoch': EPOCHS+1,
          'state_dict': model.state_dict(),
          'optimizer': optimizer.state_dict(),
          'learning_rate':lr,
        }
torch.save(state, "cow_model.pth")
</code></pre>
<p>My predictions versus ground truth look like:<br>
<strong>Prediction</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/7/b7fe97bc2513d2f86d8cd15dd368a3c1dcbfbad2.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b7fe97bc2513d2f86d8cd15dd368a3c1dcbfbad2" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/7/b7fe97bc2513d2f86d8cd15dd368a3c1dcbfbad2_2_690x263.jpeg" alt="image" data-base62-sha1="qfGQ2UD8S4eSaolMuyIMm51jusW" width="690" height="263" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/7/b7fe97bc2513d2f86d8cd15dd368a3c1dcbfbad2_2_690x263.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/7/b7fe97bc2513d2f86d8cd15dd368a3c1dcbfbad2_2_1035x394.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/7/b7fe97bc2513d2f86d8cd15dd368a3c1dcbfbad2.jpeg 2x" data-dominant-color="C3C3C3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1193×455 77.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<strong>Ground-truth</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/9/2944f31994c8663ee564cd755481a37fa70a7180.png" data-download-href="https://discuss.pytorch.org/uploads/default/2944f31994c8663ee564cd755481a37fa70a7180" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/2944f31994c8663ee564cd755481a37fa70a7180_2_689x265.png" alt="image" data-base62-sha1="5T5iBKRgkLptCxQpQ0yCNI81O6c" width="689" height="265" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/2944f31994c8663ee564cd755481a37fa70a7180_2_689x265.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/2944f31994c8663ee564cd755481a37fa70a7180_2_1033x397.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/9/2944f31994c8663ee564cd755481a37fa70a7180.png 2x" data-dominant-color="939393"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1195×459 24.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

Mukesh1729: <p>Hey Frank,</p>
<p>I tried to use the softmax in the last layer with the  MSE loss. The model predictions and ground truth look a bit similar to what the prediction was from the last post using the loss function you suggested to me.</p>
<p><strong>Model predictions</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/0/402b206d421f3d87f8d99483245d4e7affa3473c.png" data-download-href="https://discuss.pytorch.org/uploads/default/402b206d421f3d87f8d99483245d4e7affa3473c" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/0/402b206d421f3d87f8d99483245d4e7affa3473c_2_690x262.png" alt="image" data-base62-sha1="99EXj2GWFVstCkVltNo5jo9iTWk" width="690" height="262" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/0/402b206d421f3d87f8d99483245d4e7affa3473c_2_690x262.png, https://discuss.pytorch.org/uploads/default/original/3X/4/0/402b206d421f3d87f8d99483245d4e7affa3473c.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/0/402b206d421f3d87f8d99483245d4e7affa3473c.png 2x" data-dominant-color="A8A8A8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">854×325 44.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>They look quite close compared to the target but then I end up with the problem of the loss not going down. The loss you suggested is not helping the model learn the heatmaps somehow but the loss goes down in value. In the ref paper, they mention that softmax was applied as a softmax transformation together with MSE loss. Not sure what that means.<br>
Do you think there might be something that is missing or else I could try? The loss is being compared with the probability of the key points present in each heatmap to the gaussian distribution over a point in a particular location in the confidence/heat map</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/2/722f30d8e1ddd1342691d7ec6fbab328813d42b3.png" alt="image" data-base62-sha1="gi7xk35IMoYDN2lphKNQhh69U6D" width="374" height="284"></p>