Cross entropy loss ignore index

Tianyi_Zhou1: <p>Hi!</p>
<p>I am now doing a project about translation and using torch.nn.CrossEntropyLoss with torch version == 1.11.0. <a href="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html?highlight=cross%20entropy#torch.nn.CrossEntropyLoss" class="inline-onebox" rel="noopener nofollow ugc">CrossEntropyLoss — PyTorch 1.11.0 documentation</a></p>
<p>Refering to the document, I can use logits for the target instead of class indices to get the loss, so that the target/label shape will be (batchsize*sentencelength,numberofclass) in my case.</p>
<p><strong>However, the document says that I cannot use ignore_index in this case. Then, what should I do if I don’t want to calculate the loss for  token?</strong></p>
<p>Btw, I am using generated sentences data to train the mode, so it’s like a bi-level optimization end-to-end training, and the generated sentences are one-hot encoded and I cannot use argmax here, because I need to take derivative to the final loss with respect the generator model.</p>
<p>Any help will be appreciated!</p>

ptrblck: <aside class="quote no-group" data-username="Tianyi_Zhou1" data-post="1" data-topic="146784">
<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/tianyi_zhou1/48/31410_2.png" class="avatar"> Tianyi_Zhou1:</div>
<blockquote>
<p>I can use logits for the target instead of class indices to get the loss</p>
</blockquote>
</aside>
<p>NIT: you can use probabilities as the target now, not logits.</p>
<aside class="quote no-group" data-username="Tianyi_Zhou1" data-post="1" data-topic="146784">
<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/tianyi_zhou1/48/31410_2.png" class="avatar"> Tianyi_Zhou1:</div>
<blockquote>
<p>However, the document says that I cannot use ignore_index in this case. Then, what should I do if I don’t want to calculate the loss for token?</p>
</blockquote>
</aside>
<p>If “smooth” targets are passed there is no defined class index anymore, which is why <code>ignore_index</code> won’t work in this case.<br>
Instead, you could use the unreduced loss, filter out the unwanted “classes”, and reduce it afterwards as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python"># class indices
criterion = nn.CrossEntropyLoss(ignore_index=0)
input = torch.randn(3, 2, requires_grad=True)
target = torch.tensor([0, 1, 0])

loss = criterion(input, target)
print(loss)
# &gt; tensor(1.2963, grad_fn=&lt;NllLossBackward0&gt;)

# class probabilities
criterion = nn.CrossEntropyLoss(reduction='none')
target = torch.nn.functional.one_hot(target).float()

loss = criterion(input, target)
loss = loss[target[:, 0] != 1].mean()
print(loss)
# &gt; tensor(1.2963, grad_fn=&lt;MeanBackward0&gt;)
</code></pre>

Tianyi_Zhou1: <p>Hi ptrblck!<br>
Thank you so much for your answer, helps a lot!<br>
I got another short question.<br>
I am doing a bi-level optimization problem (ie. train two models). And I got some CUDA oom problem.<br>
Usually, how we train the model is like:</p>
<pre><code class="lang-auto">w_optimizer.zero_grad()
loss = CTG_loss(input_w, input_w_attn, output_w, output_w_attn, attn_idx, A, w_model)
loss.backward()
w_optimizer.step()
</code></pre>
<p>To save the memory of CUDA can i do this? In this way, after training the first model, some memory will be released.</p>
<pre><code class="lang-auto">w_optimizer.zero_grad()
loss_w = CTG_loss(input_w, input_w_attn, output_w, output_w_attn, attn_idx, A, w_model)
loss_w.backward()
w_optimizer.step()
w_optimizer.zero_grad()
torch.cuda.empty_cache()
</code></pre>

ptrblck: <p>I don’t think your second approach would avoid the OOM issue and would force PyTorch to allocate new memory in the next iteration. To avoid keeping the gradients around you could use <code>optimizer.zero_grad(set_to_none=True)</code> which could help.</p>