Loss function and LSTM dimension issues

tcsn_wty: <p>Hi all, I am writing a simple neural network using LSTM to get some understanding of NER. I understand the whole idea but got into trouble with some dimension issues, here’s the problem:</p>
<pre><code class="lang-auto">class NERModel(nn.Module):
    """
    Encoder for NER model.
    Args:
        - vocab_size: vocabulary size, integer.
        - embedding_size: embedding size, integer.
        - enc_units: hidden size of LSTM layer, integer.
        - ffc_units: hidden units of feedforward layer, integer.
        - num_labels: number of named entities. The value should be (actual_num_labels + 1),
            because zero paddings are added to the sequences.
    """
    
    def __init__(self, vocab_size, embedding_size, enc_units, ffc_units, num_labels):
        super(NERModel, self).__init__()
        # Word embedding layer.
        self.embedding = nn.Embedding(vocab_size, embedding_size)
        # LSTM layer with units of enc_units
        self.LSTM = nn.LSTM(embedding_size, enc_units, batch_first=True)
        self.dense1 = nn.Linear(enc_units, ffc_units)
        self.dense2 = nn.Linear(ffc_units, num_labels)
        
    def forward(self, x):
        """
        Args:
            - x: Input tensor of shape (batch_size, sequence_length)
        Return:
            Tensor of shape (batch_size, sequence_length, num_labels)
        """
        x = self.embedding(x)
        # after embedding: torch.Size([64, 124, 256])
        x, _ = self.LSTM(x)
        # after lstm: torch.Size([64, 124, 256])
        x = self.dense1(x)
        # after linear 1: torch.Size([64, 124, 256])
        x = self.dense2(x)
        # after linear 2: torch.Size([64, 124, 6])
        output = F.log_softmax(x, dim=1)
        # after softmax: torch.Size([64, 124, 6])
        return output

# initialize model
model = NERModel(vocab_size, embedding_size, enc_units, ffc_units, num_labels)

optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)
criterion = nn.CrossEntropyLoss()

for i, (value, label) in enumerate(train_loader):
        print(value.shape)
        optimizer.zero_grad()
        outputs = model(value)
        # outputs shape: torch.Size([64, 124, 6])
        # label shape: torch.Size([64, 124])
        loss = criterion(outputs, label)
</code></pre>
<p>Things all looked good but I got the following error reported:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-9-92530e221aaf&gt; in &lt;module&gt;()
---&gt; 14         loss = criterion(outputs, label)

~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    491             result = self._slow_forward(*input, **kwargs)
    492         else:
--&gt; 493             result = self.forward(*input, **kwargs)
    494         for hook in self._forward_hooks.values():
    495             hook_result = hook(self, input, result)

~/anaconda3/lib/python3.6/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
    940     def forward(self, input, target):
    941         return F.cross_entropy(input, target, weight=self.weight,
--&gt; 942                                ignore_index=self.ignore_index, reduction=self.reduction)
    943 
    944 

~/anaconda3/lib/python3.6/site-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
   2054     if size_average is not None or reduce is not None:
   2055         reduction = _Reduction.legacy_get_string(size_average, reduce)
-&gt; 2056     return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
   2057 
   2058 

~/anaconda3/lib/python3.6/site-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   1879         if target.size()[1:] != input.size()[2:]:
   1880             raise ValueError('Expected target size {}, got {}'.format(
-&gt; 1881                 out_size, target.size()))
   1882         input = input.contiguous().view(n, c, 1, -1)
   1883         target = target.contiguous().view(n, 1, -1)

ValueError: Expected target size (64, 6), got torch.Size([64, 124])
</code></pre>
<p>I have outputs shape torch.Size([64, 124, 6]) and label shape: torch.Size([64, 124]). It seems that loss function want me to have outputs shape torch.Size([64, 6, 124]). I don’t get the reason why it is like this, and can some one tell me how to modify it?</p>

ptrblck: <p><code>nn.CrossEntropy</code> expects a model output of the shape <code>[batch_size, nb_classes, *additional_dimensions]</code> and a target in <code>[batch_size, *additional_dimensions]</code> containing the class indices in the range <code>[0, nb_classes]</code> as explained in the <a href="https://pytorch.org/docs/stable/nn.html?highlight=crossentropyloss#torch.nn.CrossEntropyLoss">docs</a>.</p>
<p>For an output of <code>[batch_size=64, nb_classes=124, additional=6]</code>, the target should have the shape <code>[64, 6]</code> and contain values in <code>[0, 123]</code>.</p>
<p>Could you explain your use case a bit and what the dimensions in your tensors mean?</p>

tcsn_wty: <p>Thanks for the reply.</p>
<p>All commented dimensions are results by calling tensor.shape.</p>
<p>I’ll further explain the specific numbers</p>
<pre><code class="lang-auto"># print(model) result
NERModel(
  (embedding): Embedding(30290, 256)
  (LSTM): LSTM(256, 256, batch_first=True)
  (dense1): Linear(in_features=256, out_features=256, bias=True)
  (dense2): Linear(in_features=256, out_features=6, bias=True)
)

batch_size = 64
length_of_each_input = 124 words
</code></pre>
<p>As forward() function being invoked, I didn’t modify any intermediate dimensions and the output by my neural network was [64, 124, 6]. 64 is the batch size, and for each sentence with 124 words, 6 log_softmax results were calculated for each word.</p>
<p>It seems that I need to make it [64, 6, 124] so that I could call the loss function?</p>

ptrblck: <p>Yes, you should <code>permute</code> the output to be able to call <code>nn.CrossEntropyLoss</code>.<br>
I’m not sure how your <code>forward</code> is implemented, but I would recommend to check the shapes for all intermediate tensors.</p>

tcsn_wty: <p>I did output every intermediate tensor shape in forward() function, and they were printed when training. To make the code easier to read, I replaced those command with the actual results.</p>
<pre><code class="lang-auto">def forward(self, x):
        """
        Args:
            - x: Input tensor of shape (batch_size, sequence_length)
        Return:
            Tensor of shape (batch_size, sequence_length, num_labels)
        """
        x = self.embedding(x)
        # after embedding: torch.Size([64, 124, 256])
        x, _ = self.LSTM(x)
        # after lstm: torch.Size([64, 124, 256])
        x = self.dense1(x)
        # after linear 1: torch.Size([64, 124, 256])
        x = self.dense2(x)
        # after linear 2: torch.Size([64, 124, 6])
        output = F.log_softmax(x, dim=1)
        # after softmax: torch.Size([64, 124, 6])
        return output
</code></pre>
<p>As you can see, I didn’t modify the dimensions at all. Is it common that I usually need to permute the columns in order to call loss function?</p>

ptrblck: <p>The shapes look alright, if you add the discussed <code>permute</code>.<br>
However, the <code>F.log_softmax</code> operation should be applied in the class dimension, so in your case dim2 (or permute the tensor before the operation).<br>
Also, since you are using <code>F.log_softmax</code>, you should use <code>nn.NLLLoss</code> as the criterion.<br>
<code>nn.CrossEntropyLoss</code> will apply <code>F.log_softmax</code> and <code>nn.NLLLoss</code> internally.</p>

tcsn_wty: <p>Thanks! I tried <code>permute</code> just now and the error was eliminated. But if I understand correctly, you meant that I need to do the permutation before log_softmax called, am I correct?</p>
<p>As for loss function, if I’m using nn.CrossEntropyLoss, do you meant that there’s no need to apply F.log_softmax in forward() function?</p>

ptrblck: <aside class="quote no-group" data-username="tcsn_wty" data-post="7" data-topic="79291">
<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/tcsn_wty/48/11956_2.png" class="avatar"> tcsn_wty:</div>
<blockquote>
<p>But if I understand correctly, you meant that I need to do the permutation before log_softmax called, am I correct?</p>
</blockquote>
</aside>
<p>If you want to keep <code>F.log_softmax(x, dim=1)</code>, then yes.<br>
Otherwise use <code>dim=2</code>, if you want to permute the tensor afterwards.</p>
<aside class="quote no-group" data-username="tcsn_wty" data-post="7" data-topic="79291">
<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/tcsn_wty/48/11956_2.png" class="avatar"> tcsn_wty:</div>
<blockquote>
<p>As for loss function, if I’m using nn.CrossEntropyLoss, do you meant that there’s no need to apply F.log_softmax in forward() function?</p>
</blockquote>
</aside>
<p>Yes, since it will be applied twice at the moment (in your <code>forward</code> and inside <code>nn.CrossEntropyLoss</code>).</p>

vdw: <p>On a side note: You might want to add a non-linearity such as ReLU between your two linear layers. Otherwise, you don’t gain much from having two linear layers.</p>