KL loss on feature maps

username: AKASH_KUMAR
cooked: <p>Hi,</p>
<p>I am trying to minimize the distance between two feature maps using KL Divergence. The output file is of shape in both the case are : (16, 1, 8, 224, 224)</p>
<p>One map is the prediction on RGB clip and another one is the flipped version of the same clip. I’m getting a NaN loss.</p>
<pre><code class="lang-auto">output = model(data, label)  # 16, 1, 8, 224, 224
output = torch.flip(output, [4])
output = torch.mean(output, 2) # 16, 1, 224, 224
flip_op = torch.mean(flip_op, 2)
output = torch.squeeze(output, 1). # 16, 224, 224
flip_op = torch.squeeze(flip_op, 1)
# CONS_LOSS
output += 1e-7
flip_op += 1e-7
cons_loss_a = consistency_criterion(flip_op.log(), output.detach()).sum(-1).mean()
</code></pre>

username: ptrblck
cooked: <p>I’m not sure how <code>flop_op</code> is calculated, but did you make sure that the <code>log()</code> operation doesn’t produce invalid values?</p>

username: AKASH_KUMAR
cooked: <p>I was doing it the wrong way. It was not a probability. It was logits.<br>
I get flip_op the same way as output, just the data is horizontally flipped.</p>
<p>To apply KL on segmentation I tried this:</p>
<pre><code class="lang-auto">conf_consistency_criterion = torch.nn.KLDivLoss(size_average=False, reduce=False).cuda()
output = model(data, label)     # 16, 1, 8, 224, 224  
flip_op = model(flipped_data, label)   # 16, 1, 8, 224, 224
output = torch.flip(output, [4])
output = torch.mean(output, 2) # 16, 1, 224, 224
flip_op = torch.mean(flip_op, 2) # 16, 1, 224, 224
output = torch.squeeze(output, 1) # 16, 224, 224
flip_op = torch.squeeze(flip_op, 1) # 16, 224, 224
# CONS_LOSS
output += 1e-7
flip_op += 1e-7
cons_loss_a = conf_consistency_criterion(F.logsigmoid(output), F.sigmoid(flip_op)).sum(-1).mean()
</code></pre>
<p>I’m getting a negative output. So, this is wrong somewhere. KL is a measure of how close two probability distributions.</p>
<p>I also tried this:</p>
<pre><code class="lang-auto">conf_consistency_criterion = torch.nn.KLDivLoss(size_average=False, reduce=False).cuda()
output = model(data, label)     # 16, 1, 8, 224, 224  
flip_op = model(flipped_data, label)   # 16, 1, 8, 224, 224
output = torch.flip(output, [4])    
# CONS_LOSS
output += 1e-7
flip_op += 1e-7
cons_loss_a = conf_consistency_criterion(F.logsigmoid(output), F.sigmoid(flip_op)).sum(-1).mean()
</code></pre>
<p>This gives me positive value. However, the value suddenly drops from 20 to 0.006 in the second epoch. I don’t get what’s the right approach to do so?</p>

username: ptrblck
cooked: <p>I’m not getting a negative loss in any of the two approaches:</p>
<pre><code class="lang-python">conf_consistency_criterion = torch.nn.KLDivLoss(size_average=False, reduce=False).cuda()
output = torch.randn(16, 1, 8, 224, 224  )
flip_op = torch.randn(16, 1, 8, 224, 224)
output = torch.flip(output, [4])
output = torch.mean(output, 2) # 16, 1, 224, 224
flip_op = torch.mean(flip_op, 2) # 16, 1, 224, 224
output = torch.squeeze(output, 1) # 16, 224, 224
flip_op = torch.squeeze(flip_op, 1) # 16, 224, 224
# CONS_LOSS
output += 1e-7
flip_op += 1e-7
cons_loss_a = conf_consistency_criterion(F.logsigmoid(output), torch.sigmoid(flip_op)).sum(-1).mean()
print(cons_loss_a)


output = torch.randn(16, 1, 8, 224, 224)
flip_op = torch.randn(16, 1, 8, 224, 224)
output = torch.flip(output, [4])    
# CONS_LOSS
output += 1e-7
flip_op += 1e-7
cons_loss_a = conf_consistency_criterion(F.logsigmoid(output), F.sigmoid(flip_op)).sum(-1).mean()
print(cons_loss_a)
</code></pre>
<p>so don’t know what might be causing this issue.</p>

username: AKASH_KUMAR
cooked: <p>So, either approach can be used to compute the JSD loss?<br>
Max min values for both of my segmentation maps range between 2 to -2.<br>
My ground truth map is binary 0/1.</p>

username: ptrblck
cooked: <p>Yes, I think so. Based on the <a href="https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html#torch.nn.KLDivLoss">docs</a> the inputs should be given as log-probabilities and the targets as probabilities by default, which seems to be the case for you.<br>
I’m not familiar with your use case so cannot comment, if <code>(log)sigmoid</code> or <code>(log)softmax</code> should be used.</p>

username: AKASH_KUMAR
cooked: <p>I have a binary output. So, sigmoid is the right choice, isn’t that so?</p>