How to properly use ignore_index when training a transformer

Theo: <p>I’m trying to train a very simple transformer which is supposed to translate sentences to another language and I’m having trouble getting rid of padding index when training my model.</p>
<p>Here is my training loop(simplified for clarity):</p>
<pre><code class="lang-auto">for source, target in data: # Source and target Shape [ max_len ]
    source, target = source.to(device), target .to(device)
    optimizer.zero_grad()
    logits = model(source, target ) # Shape [ max_len, vocab_size ]
    loss = criterion(logits, labels) # criterion is CrossEntropyLoss
    loss.backward()
    optimizer.step()
</code></pre>
<p>After training, I then do<br>
<code>logits = model(source, target)</code><br>
and<br>
<code>torch.argmax(logits, dim=1)</code><br>
which gives me the ids of words of the translated sentence.</p>
<p>An example source sentence:</p>
<pre><code class="lang-auto">tensor([    2,   214,   544,  1392,   306,   708,   168,   743,   103,   145,
           839, 16289, 12140,    15,     3,     1,     1,     1,     1,     1])
</code></pre>
<p>Target sentence:</p>
<pre><code class="lang-auto">tensor([   2,  217,  273,   13, 5008, 3650,  192,  479,  581, 1788, 5743,   15,
            3,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1,    1])
</code></pre>
<p>My labels are of shape [ max_len, vocab_size ], and every word is given a probability of .0 except the correct word which is given 1.0. For example, labeling of a random 5 word sentence in a language with a vocab_size of 10 words would look like this:</p>
<pre><code class="lang-auto">tensor([[     .0,     .0,    1.0,     .0,     .0,     .0,     .0,     .0,     .0,     .0]
        [     .0,     .0,     .0,    1.0,     .0,     .0,     .0,     .0,     .0,     .0]
        [     .0,     .0,     .0,     .0,     .0,    1.0,     .0,     .0,     .0,     .0]
        [     .0,     .0,     .0,     .0,     .0,     .0,     .0,     .0,     .0,    1.0]
        [     .0,     .0,     .0,    1.0,     .0,     .0,     .0,     .0,     .0,     .0]])
</code></pre>
<p>CrossEntropyLoss then takes the output of my model and the labels of the “target” sentence as input.</p>
<p>The problem is this: Since there are lots of paddings(padding index is 1 in my case), the model eventually tries to translate the sentence into only padding words because they are the majority, thus resulting in a smaller loss.</p>
<p>I tried using the “ignore_index” feature of the CrossEntropyLoss() function like this<br>
<code>CrossEntropyLoss(ignore_index=1)</code> but it says:</p>
<pre><code class="lang-auto">RuntimeError: ignore_index is not supported for floating point target
</code></pre>
<p>How can I solve this problem?</p>

ptrblck: <p>As the error message explains, the <code>ignore_index</code> argument is invalid for (the new) floating point targets.<br>
Assuming you don’t want to use “soft” target values and could use class indices, you could transform the current target to class indices via:</p>
<pre><code class="lang-python">target = torch.tensor([[     .0,     .0,    1.0,     .0,     .0,     .0,     .0,     .0,     .0,     .0],
                       [     .0,     .0,     .0,    1.0,     .0,     .0,     .0,     .0,     .0,     .0],
                       [     .0,     .0,     .0,     .0,     .0,    1.0,     .0,     .0,     .0,     .0],
                       [     .0,     .0,     .0,     .0,     .0,     .0,     .0,     .0,     .0,    1.0],
                       [     .0,     .0,     .0,    1.0,     .0,     .0,     .0,     .0,     .0,     .0]])
target = target.argmax(dim=1)
print(target)
# tensor([2, 3, 5, 9, 3])
</code></pre>
<p>which would not only save memory but would also allow you to use <code>ignore_index</code> properly.</p>

Theo: <p>Thanks for the reply, my loss function now works properly.</p>
<p>I wanted to ask one more thing: does <code>CrossEntropyLoss(logits, target)</code> realize that when logits is of shape <code>[ max_len, vocab_size ]</code>, the <code>max_len</code> dimension does not represent batches but rather the number of words in my sentence?</p>
<p>In the documentation, it expects logits to be of shape <code>(N, C)</code>, however, here <code>N</code> represents number of batches, which I don’t have because I work with only one sentence at a time.</p>

ptrblck: <p>No, the loss function doesn’t understand any difference between number of samples in the batch dimension or number of sentences and will treat <code>dim0</code> as the dimension containing different samples.</p>

Theo: <p>Thanks for the reply, does that mean that there is no way for me to train my model with batches? My input then would have to be of shape [ batch_size, max_len, vocab_size ] which is from my understanding not an acceptable input shape for the loss function. Is there a way around this?</p>

ptrblck: <p>No, sorry this is not what I meant.<br>
In my previous answer I assumed you asked if the loss function would understand the difference between a model output in the shape <code>[batch_size, nb_classes]</code> and <code>[max_len, nb_classes]</code>, which is not the case.<br>
<code>nn.CrossEntropyLoss</code> accepts multi-dimensional data, in your case temporal data in the shape <code>[batch_size, nb_classes, sequence_length]</code> as the model output and expects targets in the shape <code>[batch_size, sequence_length]</code> containing class indices in the range <code>[0, nb_classes-1]</code>.<br>
You might thus need to <code>.permute</code> your model output to create the expected shape.</p>