Train/validation loss not decreasing

Mukesh1729: <p>Hi, I am taking the output from my final convolutional transpose layer into a softmax layer and then trying to measure the mse loss with my target. Problem is that my loss is doesn’t decrease and is stuck around the same point. I figured the problem is using the softmax in the last layer.</p>
<p>I need the softmax layer in the last layer because I want to measure the probabilities. Any suggestion on how to overcome this problem would be helpful. Thanks</p>
<p>Ex:<br>
input  = torch.randn(1, 16, 2, 256, 256)<br>
model output  = torch.randn(1, 16, 1, 256, 256)<br>
target = torch.randn(1, 16, 1, 256, 256)<br>
criterion = mse()</p>

Azerus: <p>Hello,</p>
<p>If you want to your network to output a probability distribution, it’s better to use a cross entropy loss (in case you have one_hot target, as for classification, or BCE loss in case of 2 classes) or a KL-divergence loss (in case you have a target distribution). Those losses work generally better than mse in such cases (there are theoretical justification of this claim). Little advice, if you want to use cross entropy loss, do not insert a softmax at the end of your model, <a href="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#" rel="noopener nofollow ugc">CrossEntropyLoss</a> implemented on pytorch works directly with input logits for a better numerical precision and stability.</p>
<p>Hope it helps,<br>
Thomas</p>

Mukesh1729: <p>Hey Thomas,<br>
thanks for getting back, I am doing a regression-based pose estimation problem. From the paper, I have found they use a mse loss function with a softmax transformation to predict the key points and then compute loss. My only problem is that the loss is stuck when the softmax is present. I want to know if there is some trick when using softmax with mse loss to avoid the loss problem</p>
<p>My model can make predictions but the loss doesnt improve</p>
<pre><code class="lang-auto">class model(nn.Module):

    def __init__(self):
        super(model, 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):
        """PUTTING THE MODEL TOGETHER"""
        encoder = self.encoder(image)
        decoder = self.decoder(encoder)
        return torch.nn.Softmax(dim=1)(decoder)


model = model()
if torch.cuda.is_available():
  input = torch.rand(1,3, 2,  256, 256).cuda()
  model = model.cuda()
  print("Output image shape:",model(input).shape)
  
else:
  summary(model,
          input_size=(3, 2, 200, 200),
          batch_size=1
          )
print("************************************************")
# torch.cuda.memory_summary(device=None, abbreviated=False)
</code></pre>
<p>From the paper:</p>
<ul>
<li>
<p>The last convolutional layer of the decoder is followed by a<br>
transposed-convolutional layer that uses a linear activation followed by a Softmax transformation, resulting in an output of 17 confidence maps, one per keypoint.</p>
</li>
<li>
<p>The loss function calculated the Mean Squared Error (MSE) per pixel per map between the predicted confidence maps and the ground-truth confidence maps from the<br>
samples in the batch.</p>
</li>
</ul>

Azerus: <p>Mmmh, I don’t know such trick. Could you send a link to the paper? I’ll have a look (probably tomorrow, I won’t have much time today).</p>

Mukesh1729: <p>Yeah sure,</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://arxiv.org/abs/2104.08029">
  <header class="source">
      <img src="https://static.arxiv.org/static/browse/0.3.2.8/images/icons/favicon.ico" class="site-icon" width="32" height="32">

      <a href="https://arxiv.org/abs/2104.08029" target="_blank" rel="noopener nofollow ugc">arXiv.org</a>
  </header>

  <article class="onebox-body">
    <img src="https://static.arxiv.org/icons/twitter/arxiv-logo-twitter-square.png" class="thumbnail onebox-avatar" width="500" height="500">

<h3><a href="https://arxiv.org/abs/2104.08029" target="_blank" rel="noopener nofollow ugc">T-LEAP: occlusion-robust pose estimation of walking cows using temporal...</a></h3>

  <p>As herd size on dairy farms continue to increase, automatic health monitoring
of cows has gained in interest. Lameness, a prevalent health disorder in dairy
cows, is commonly detected by analyzing the gait of cows. A cow's gait can be
tracked in...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Azerus: <p>Ok, I had the time to have a quick look. I hope I’m not wrong, but from what I understood the softmax layer is used per map, and not per pixel. If I understood right looking at the generated confidence-maps page 6, the softmax is applied on the spatial dimension, not across the channels, which mean you want a spatial distribution and not a class distribution.</p>
<p>So If you follow the pytorch convention for the dmension of your images (Batch, channels, x, y), then you should apply your softmax as:</p>
<pre><code class="lang-auto">return torch.nn.Softmax(dim=(2,3))(decoder)
</code></pre>
<p>I hope it should be right like that.</p>

Azerus: <p>Ah, sorry, softmax dim option cannot take tuple but only int, so you have to flatten your image before computing softmax, something like this should do the trick:</p>
<pre><code class="lang-auto">decoder_shape = decoder.shape
flatten_decoder = decoder.view(decoder_shape[0], decoder_shape[1], -1)
faltten_heat_map = torch.nn.Softmax(dim=2)(flatten_decoder )
return faltten_heat_map.view(decoder_shape)
</code></pre>

Mukesh1729: <p>Hey Thomas,<br>
If I do this then i need to reshape my ground truth heatmap too</p>
<pre><code class="lang-auto">decoder_shape = decoder.shape
flatten_decoder = decoder.view(decoder_shape[0], decoder_shape[1], -1)
faltten_heat_map = torch.nn.Softmax(dim=2)(flatten_decoder )
return faltten_heat_map.view(decoder_shape)
</code></pre>

Azerus: <p>You do not have to reshape your target because of the last view <code>faltten_heat_map.view(decoder_shape)</code> that reshape correctly the predictions</p>

Mukesh1729: <p>Hey Thomas, I tried what you suggested and the outcome looks like the given plot</p>
<pre><code class="lang-auto">class T_LEAP(nn.Module):
    """T_LEAP ARCHITECTURE"""

    def __init__(self):
        super(T_LEAP, 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):
        """PUTTING THE MODEL TOGETHER"""
        encoder = self.encoder(image)
        decoder = self.decoder(encoder)
        decoder_shape = decoder.shape
        flatten_decoder = decoder.view(decoder_shape[0], decoder_shape[1], -1)
        flatten_heat_map = torch.nn.Softmax(dim=2)(flatten_decoder )
        return flatten_heat_map.view(decoder_shape)


model = T_LEAP()
if torch.cuda.is_available():
  input = torch.rand(1,3, 2,  256, 256).cuda()
  model = model.cuda()
  print("Output image shape:",model(input).shape)
  
else:
  summary(model,
          input_size=(3, 2, 200, 200),
          batch_size=1
          )
print("************************************************")
# torch.cuda.memory_summary(device=None, abbreviated=False)
</code></pre>
<p>Sample training loop:</p>
<pre><code class="lang-auto">#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)
model = model.to(DEVICE)


# def softXEnt(input, target):
#     logprobs = torch.nn.functional.log_softmax (input, dim = 1)
#     return  -(target * logprobs).mean()

tr_losses =[]
val_losses=[]
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()

image   = torch.rand(1,3,2,160, 160).to(DEVICE)
heatmap = torch.rand(1,16,1,160,160).to(DEVICE)

for idx in range(EPOCHS):
  
  model.train()
  
  # Get images and transfer to GPU
  image   = image
  heatmap = heatmap
  
  optimizer.zero_grad()
  output = model(image)
  loss   = criterion(output,heatmap)
  loss.backward()
  optimizer.step()
  
  curr_trloss = loss.item()
  # Print losses 
  print(f"Training loss {idx}:",curr_trloss/1)
        

# Evaluation loop
  model.eval()
  with torch.no_grad():
    # Get images and transfer to GPU
    image   = image
    heatmap = heatmap
    
    output = model(image)
    loss   = criterion(output,heatmap)
    curr_valoss = loss.item()
    
  print(f"Validation loss {idx}:",curr_valoss/1)
  print(" ")
  tr_losses.append(curr_trloss/1)
  val_losses.append(curr_valoss/1)
    
print("Training and evalutation is now complete")

loss_plot(EPOCHS,tr_losses, val_losses)
</code></pre>
<p><strong>Output</strong><br>
yellow - validation curve (sorry about the legends)<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/4/6/463fc3b746dd2eb8ce830d949d6a32c29b1d4ce7.png" alt="image" data-base62-sha1="a1s2lcX4dMz07MSySVkBvjG79mD" width="329" height="233"><br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6d74024c9b74118b44e880f867579e4863a1d1d4.png" alt="image" data-base62-sha1="fCgAJ5DBIdAyBEGAet8aUFVAhWQ" width="300" height="479"></p>

Azerus: <p>If i understood right, you are doing a dummy training with a constant image and heat map ? I see at least one issue, the heat map is not normalized. When defining your dummy heat map do this:</p>
<pre><code class="lang-auto">heatmap = torch.rand(1,16,1,160,160).to(DEVICE)
heatmap = heatmap/torch.sum(dim=(2,3,4), keepdim=True)
</code></pre>
<p>But there is maybe something else that disturb the training… Still looking</p>

Mukesh1729: <aside class="quote no-group" data-username="Mukesh1729" data-post="11" data-topic="137889">
<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>
<pre><code class="lang-auto"># def softXEnt(input, target):
#     logprobs = torch.nn.functional.log_softmax (input, dim = 1)
#     return  -(target * logprobs).mean()
</code></pre>
</blockquote>
</aside>
<p>Someone from the forums suggested me to remove the Softmax and mse loss and try this instead:</p><aside class="quote quote-modified" data-post="8" data-topic="137473">
  <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">
    <a href="https://discuss.pytorch.org/t/compute-mse-loss-with-softmax/137473/8">Compute mse_loss() with softmax()</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    Hey Frank, 
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. 
Model predictions 
[image] 
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 mentio…
  </blockquote>
</aside>

<pre><code class="lang-auto">def softXEnt(input, target):
     logprobs = torch.nn.functional.log_softmax (input, dim = 1)
     return  -(target * logprobs).mean()
</code></pre>
<p>In this case, the loss decreases but the model learns poorly as i can see from the predictions</p>

Mukesh1729: <ul>
<li>I do normalize my heatmaps while training though.</li>
<li>If you remove the softmax loss the loss will start go down though. In the paper they mention about using a linear activation before the softmax, I am assuming it is doing nothing in PyTorch</li>
</ul>
<p>throws back an error</p>
<pre><code class="lang-auto">heatmap = heatmap/torch.sum(dim=(2,3,4), keepdim=True)

</code></pre>
<p>Sorry to keep adding info along the way. I thought it could a helpful reference to rule out or consider new possibilities</p>

Azerus: <p>Yes, working on log probabilities is not a bad idea, but be cautious, in the code snipped you send me the log softmax is applied on channels (dim=1) instead of spatial dimensions.</p>
<p>Yes, sorry, I do some errors when writing code too quickly ^^ I indeed wanted to say this:</p>
<pre><code class="lang-auto">heatmap = torch.rand(1,16,1,160,160).to(DEVICE)
heatmap = heatmap/torch.sum(heatmap , dim=(2,3,4), keepdim=True)
</code></pre>
<p>And yes, when someone say “linear activation”, usually it means no activation function at output of the layer (as classical layers are linear).</p>

Mukesh1729: <p>Hey Thomas,<br>
I am still encountering the same issue. Do you think my code is okay so far in comparison with the paper?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/9/a9232ec4014eea1501d8f77d4885c3e062a22e87.png" data-download-href="https://discuss.pytorch.org/uploads/default/a9232ec4014eea1501d8f77d4885c3e062a22e87" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/9/a9232ec4014eea1501d8f77d4885c3e062a22e87.png" alt="image" data-base62-sha1="o8g4p6n3BRzH1M4hwhXGIibuRhl" width="233" height="500" data-dominant-color="6A6470"><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">353×757 29.7 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>

Azerus: <p>At least your training curve seems more consistent now, which is a good news! I don’t understand the validation curve, when you do this (your snipped of code, no modification from my side here):</p>
<pre><code class="lang-auto">for idx in range(EPOCHS):
  
  model.train()
  
  # Get images and transfer to GPU
  image   = image
  heatmap = heatmap
  
  optimizer.zero_grad()
  output = model(image)
  loss   = criterion(output,heatmap)
  loss.backward()
  optimizer.step()
  
  curr_trloss = loss.item()
  # Print losses 
  print(f"Training loss {idx}:",curr_trloss/1)
        

# Evaluation loop
  model.eval()
  with torch.no_grad():
    # Get images and transfer to GPU
    image   = image
    heatmap = heatmap
    
    output = model(image)
    loss   = criterion(output,heatmap)
    curr_valoss = loss.item()
    
  print(f"Validation loss {idx}:",curr_valoss/1)
  print(" ")
  tr_losses.append(curr_trloss/1)
  val_losses.append(curr_valoss/1)
</code></pre>
<p>The during both the training phase and evaluation phase you are using the same input image and target heatmap. So I don’t get why the validation loss is not equal to the training loss … Are you using this snipped, or running something different?</p>

Mukesh1729: <p>I am using the same snippet of code you’ve mentioned in your post</p>

Azerus: <p>Well the only difference I see in training phase VS validation phase is that the batch norm will behave differently in train mode VS eval mode. Could you try running again your code but commenting (just for this test) all the batch norm layers in your model?</p>

Mukesh1729: <p>I have commented out all of the batchnorm layers from the model and the output looks like:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/818b6768c520d44cf07a6128c54698eb4e70aa53.png" data-download-href="https://discuss.pytorch.org/uploads/default/818b6768c520d44cf07a6128c54698eb4e70aa53" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/818b6768c520d44cf07a6128c54698eb4e70aa53_2_223x500.png" alt="image" data-base62-sha1="iu0grEbTVw08ZGfKJVq5wwohuqn" width="223" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/818b6768c520d44cf07a6128c54698eb4e70aa53_2_223x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/818b6768c520d44cf07a6128c54698eb4e70aa53_2_334x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/1/818b6768c520d44cf07a6128c54698eb4e70aa53.png 2x" data-dominant-color="6B6571"><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">335×751 38.6 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><strong>output with batchnorm commented out and output layer is followed by just torch.nn.Softmax(dim=1)</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c3cc9a0bc35164fc6dc86ee15a0cb8b298005d7.png" data-download-href="https://discuss.pytorch.org/uploads/default/0c3cc9a0bc35164fc6dc86ee15a0cb8b298005d7" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/c/0c3cc9a0bc35164fc6dc86ee15a0cb8b298005d7_2_260x500.png" alt="image" data-base62-sha1="1KfXVdeO4PwA87kNDWJcPSMqnu7" width="260" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/c/0c3cc9a0bc35164fc6dc86ee15a0cb8b298005d7_2_260x500.png, https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c3cc9a0bc35164fc6dc86ee15a0cb8b298005d7.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/c/0c3cc9a0bc35164fc6dc86ee15a0cb8b298005d7.png 2x" data-dominant-color="706B76"><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">363×698 34.6 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>

Azerus: <p>Ok, so it was the batch norm. Now you’re validation curve and training curves are consistent with respect to each other at least.<br>
I’m just not sure to get the difference between the two experiments. So, in the first one you used the flatten trick to apply softmax on the spatial dimensions, and on the second one you applied the softmax across the channels (dim=1), is that correct?</p>