Reimplementing CrossEntropyLoss leads to unexpected gradient vanishing

Crow: <p>My task is to make multi labels classification which I choose CrossEntropyLoss to do so. However, I have to reimplement a CELoss function as the pytorch built-in CELoss only accept a single label tensor as target (shape [batch, ]). Here is the loss part of my code:</p>
<pre><code class="lang-auto"># using built-in loss function
'''  
labels[:, 1:] = 0  # to debug, I simplify the task to single label classification.
labels = labels.sum(-1)
logits = logits.view(-1, class_num)
labels = labels.view(-1)
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits, labels)
'''

# my implementation
bs, label_num = labels.size()
labels[:, 1:] = 0
log_softmax_output = torch.log(F.softmax(logits, dim=-1))
log_softmax_output = log_softmax_output.view(-1, class_num)

w = torch.arange(label_num, dtype=logits.dtype, device=logits.device).view(1, label_num).tile([bs, 1])
nonzero_label_num = (labels &gt; 0).to(logits.dtype).sum(dim=-1, keepdim=True)
w = ((nonzero_label_num - w) / nonzero_label_num) * (labels &gt; 0).to(logits.dtype)

labels = labels.view(-1, label_num)
w = w.view(-1, label_num)
w = w[labels &gt; 0] #[batchsize * nonzero_label_num]

relevant_logits = torch.gather(log_softmax_output, dim=-1, index=labels)
relevant_logits = relevant_logits[labels &gt; 0] * w
loss2 = -relevant_logits.mean()
#print(loss, loss2)
</code></pre>
<p>The labels is a list zero-padding to a fixed <strong>label_num</strong> and the <strong>nonzero_label_num</strong> denotes actual class num of this sample. The label importance decrease from index 0 to <strong>nonzero_label_num</strong> which I use <strong>w</strong> to represent the multi labels importance. The problem becomes a common single label classification when I set <code>label[:, 1:] = 0</code>.</p>
<p>Under the single label prediction condition (to debug the code), I ran the code and the output shows that loss and loss2 are equal.</p>
<p>But the strange thing appears when the training procedure comes to a specific step. The gradients stop changing in the second implementation as well as the  cosine LR_scheduler,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/9/f946914e52d712541b13b1c11d23503c3cb54d44.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/f946914e52d712541b13b1c11d23503c3cb54d44" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f946914e52d712541b13b1c11d23503c3cb54d44_2_690x144.jpeg" alt="image" data-base62-sha1="zzc4ZyFkTAwO3dJIGbLa1AgCqSU" width="690" height="144" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f946914e52d712541b13b1c11d23503c3cb54d44_2_690x144.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f946914e52d712541b13b1c11d23503c3cb54d44_2_1035x216.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/9/f946914e52d712541b13b1c11d23503c3cb54d44_2_1380x288.jpeg 2x" data-dominant-color="F6F3F3"><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">2780×584 215 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>
while the built-in implementation works fine.</p>
<p>So what’s the problem? Can any one help? Thanks so much.</p>

Crow: <p>the stdout:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b5964aad506206e7dd3a5e88f25276ca005e6003.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/b5964aad506206e7dd3a5e88f25276ca005e6003" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5964aad506206e7dd3a5e88f25276ca005e6003_2_690x68.jpeg" alt="image" data-base62-sha1="pUoq4YliTi0wrZMxgWJ3G0aZLY7" width="690" height="68" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5964aad506206e7dd3a5e88f25276ca005e6003_2_690x68.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5964aad506206e7dd3a5e88f25276ca005e6003_2_1035x102.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/5/b5964aad506206e7dd3a5e88f25276ca005e6003_2_1380x136.jpeg 2x" data-dominant-color="344149"><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">1720×170 135 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>

Crow: <p>built-in CELoss:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd1c2e79eb0aa099ba76b581cd7b6be4851d9c67.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bd1c2e79eb0aa099ba76b581cd7b6be4851d9c67" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd1c2e79eb0aa099ba76b581cd7b6be4851d9c67_2_690x142.jpeg" alt="image" data-base62-sha1="qYWCtAYOhJqDFowkJOFMRs5lrn1" width="690" height="142" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd1c2e79eb0aa099ba76b581cd7b6be4851d9c67_2_690x142.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd1c2e79eb0aa099ba76b581cd7b6be4851d9c67_2_1035x213.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/d/bd1c2e79eb0aa099ba76b581cd7b6be4851d9c67_2_1380x284.jpeg 2x" data-dominant-color="EFF2F5"><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">2766×570 218 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>

ptrblck: <p>Replace <code>log_softmax_output = torch.log(F.softmax(logits, dim=-1))</code> with <code>F.log_softmax</code> for more numerical stability.<br>
Besides that you could use <code>nn.BCEWithLogitsLoss</code> for a multi-label classification where each sample belongs to zero, one, or multiple classes.</p>

Crow: <p>Problem solved.  Many thanks!!!  <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"> <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"></p>