Creating a custom BCE with logit loss function

chinmay5: <p>I want to create a custom loss function for multi-label classification. The idea is to weigh the positive and negative labels differently. For this, I am making use of this custom code implementation.</p>
<pre><code class="lang-auto">class WeightedBCEWithLogitLoss(nn.Module):
    def __init__(self, pos_weight, neg_weight):
        super(WeightedBCEWithLogitLoss, self).__init__()
        self.register_buffer('neg_weight', neg_weight)
        self.register_buffer('pos_weight', pos_weight)

    def forward(self, input, target):
        assert input.shape == target.shape, "The loss function received invalid input shapes"
        y_hat = torch.sigmoid(input + 1e-8)
        loss = -1.0 * (self.pos_weight * target * torch.log(y_hat + 1e-6) + self.neg_weight * (1 - target) * torch.log(1 - y_hat + 1e-6))
        # Account for 0 times inf which leads to nan
        loss[torch.isnan(loss)] = 0
        # We average across each of the extra attribute dimensions to generalize it
        loss = loss.mean(dim=1)
        # We use mean reduction for our task
        return loss.mean()
</code></pre>
<p>I started getting <code>nan</code> values which I realized happened because of 0 times inf multiplication. I handled it as shown in the figure. Next, I again saw getting <code>inf</code> as the error value and corrected it by adding 1e-6 to the log (I tried with 1e-8 but that still gave me inf error value).</p>
<p>It would be great if someone can take a look and suggest further improvements and rectify any more bugs visible here.</p>

KFrank: <p>Hi Chinmay!</p>
<aside class="quote no-group" data-username="chinmay5" data-post="1" data-topic="122459" 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/chinmay5/48/12682_2.png" class="avatar"> chinmay5:</div>
<blockquote>
<p>I want to create a custom loss function for multi-label classification. The idea is to weigh the positive and negative labels differently.</p>
</blockquote>
</aside>
<p>Do be aware that pytorch’s <code>BCEWithLogitsLoss</code> supports a<br>
<code>pos_weight</code> constructor argument that will do what you want.<br>
So unless this is a learning exercise, you should simply use<br>
<code>BCEWithLogitsLoss</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        y_hat = torch.sigmoid(input + 1e-8)
</code></pre>
</blockquote>
</aside>
<p>The <code>1e-8</code> doesn’t do anything useful here.  <code>sigmoid()</code> is very<br>
well behaved (and equal to <code>0.5</code>) when its argument is equal to<br>
zero.  So there’s no need or benefit in trying to move the argument<br>
a little bit away from zero.  (Furthermore, <code>-1.e-8</code> is a perfectly<br>
valid logit and argument to your loss function and your “fix” just<br>
moves it to zero – not that anything bad happens at zero.)</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">        loss = -1.0 * (self.pos_weight * target * torch.log(y_hat + 1e-6) + self.neg_weight * (1 - target) * torch.log(1 - y_hat + 1e-6))
</code></pre>
</blockquote>
</aside>
<p>Here you apply <code>log()</code> to <code>sigmoid()</code>.  This is a source of numerical<br>
instability.  You should use the <a href="https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations" rel="noopener nofollow ugc">log-sum-exp</a> “trick” to compute<br>
<code>log (sigmoid())</code>.  (This is what pytorch’s <code>BCEWithLogitsLoss</code><br>
does internally.)</p>
<aside class="quote no-group">
<blockquote>
<p>I started getting <code>nan</code> values<br>
…<br>
Next, I again saw getting <code>inf</code></p>
</blockquote>
</aside>
<p>With the code you posted, I don’t see why you would be getting<br>
<code>nan</code>s or <code>inf</code>s.  The <code>1.e-6</code> that you add to your <code>log()</code> functions<br>
should protect against that.</p>
<p>Best.</p>
<p>K. Frank</p>