nn.LSTM with nn.NLLLoss: confusion between batch size and sequence length

username: yaskovdev
cooked: <p>Good evening!</p>
<p>I am trying to follow the <a href="https://pytorch.org/tutorials/beginner/nlp/sequence_models_tutorial.html" rel="noopener nofollow ugc">Sequence Models and Long Short-Term Memory Networks</a> tutorial. I’m confused about how to use batches with <code>nn.LSTM</code>.</p>
<p>The tutorial says at the beginning: “The first axis is the sequence itself, the second indexes instances in the mini-batch, and the third indexes elements of the input. We haven’t discussed mini-batching, so let’s just ignore that and assume we will always have just 1 dimension on the second axis.” This suggests that we will be dealing with batches of size 1.</p>
<p>But then the tutorial instanciates the <code>nn.NLLLoss</code> function and in the first iteration passes <code>tag_scores</code> of shape <code>(5, 3)</code> and <code>targets</code> of shape <code>(5)</code> to it. According to <a href="https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html" rel="noopener nofollow ugc">the docs of the loss function</a>, the first dimension is the size of the mini-batch. Which suggests that in the tutorial the batches are of size 5 now.</p>
<p>Which is correct? Isn’t 5 the sequence length (the length of the “the dog ate the apple” sentence), not the mini-batch size?<br>
If instead of 1 sequence of length 5 I had, say, a mini-batch of 32 sequences of length 5, what would be the shape of the <code>nn.NLLLoss</code> input then? <code>(160, 3)</code>, <code>(160)</code>?</p>
<p>Thank you.</p>

username: vdw
cooked: <p>Yes, the sequence length is 5. The 3 stems from the 3 probabilities that you get for each token since you have 3 possible tags (in this simplified) example here. For example, if you replace</p>
<pre><code class="lang-auto">tag_scores = F.log_softmax(tag_space, dim=1)
</code></pre>
<p>with</p>
<pre><code class="lang-auto">tag_scores = F.softmax(tag_space, dim=1)
</code></pre>
<p>the 3 probabilities in each of the 5 rows should sum up to 1. The value in <code>targets</code> are 0, 1, or 2, reflecting with class/probability is the correct one for each token.</p>
<p>Side note: I’m note sure, but extending this tutorial to mini-batches might be tricky. This tutorial uses a lot of <code>view()</code> commands that might cause <a href="https://discuss.pytorch.org/t/for-beginners-do-not-use-view-or-reshape-to-swap-dimensions-of-tensors/75524">problems</a>.</p>

username: yaskovdev
cooked: <p>Thank you for the reply, <a class="mention" href="/u/vdw">@vdw</a>!</p>
<p>So, if I understand correctly, if the output of my <code>nn.LSTM</code> layer has a shape of <code>(BATCH_SIZE, SEQ_LENGTH, HIDDEN_DIM)</code>, then I the arguments of the <code>nn.NLLLoss</code> instance (its input and target) should have the next shape: <code>(BATCH_SIZE * SEQ_LENGTH, HIDDEN_DIM)</code> and <code>(BATCH_SIZE * SEQ_LENGTH)</code>.</p>
<p>At least when I <a href="https://github.com/jshaipuka/ai-sandbox/blob/master/music-generation/train.py" rel="noopener nofollow ugc">try calling the <code>nn.NLLLoss</code> instance this way</a>, it works as expected.</p>

username: vdw
cooked: <p>You need to be careful! Have a look how the <code>nn.LSTM</code> layer has been defined:</p>
<pre><code class="lang-auto">self.lstm = nn.LSTM(embedding_dim, hidden_dim)
</code></pre>
<p>This means it uses the default input parameter <code>batch_first=False</code>. If you check the <a href="https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html" rel="noopener nofollow ugc">docs</a> for <code>nn.LSTM</code>, the shape of <code>lstm_out</code> is <code>(seq_len, batch_size, hidden_size)</code> – not this also assume that <code>bidirectional=False</code>.</p>
<p>Before you push this through <code>self.hidden2tag</code>, you need to get it into the shape of <code>(batch_size, seq_len, hidden_size)</code>. The tutorial uses <code>.view()</code> for that which I would highly discourage. I would recommend</p>
<pre><code class="lang-auto">lstm_out = lstm_out.transpose(1,0)
</code></pre>
<p>or</p>
<pre><code class="lang-auto">lstm_out = lstm_out.permute(1,0,2)
</code></pre>
<p>Then, after the call on that “corrected” <code>lstm_out</code></p>
<pre><code class="lang-auto">tag_space = self.hidden2tag(lstm_out)
</code></pre>
<p><code>tag_space</code> will have a shape of <code>(batch_size, seq_len, 3)</code>, which is what you want: for each sequence in your batch and each word in a sequence, you want the 3 probabilities.</p>
<p>In the general case – i.e., when using batches with more than 1 sequence – the shape of targets then is <code>(batch_size, seq_len)</code>, where the values at each position of the tensor are 0, 1, or 2, indicating the correct class / label.</p>
<p>IMPORTANT: Since this tutorial uses only batches of size 1, things are simpler and the <code>.view()</code> command probably works. For batches of size &gt; 1, I’m doubt this code will still work properly.</p>
<p>I have a complete example for an RNN-based NER tagger, which is the same setup but predicting NER tags for each token/word instead of POS tags. Here’s the <a href="https://github.com/chrisvdweth/nus-cs5246/blob/master/notebooks/Named%20Entity%20Recognition%20(RNN).ipynb" rel="noopener nofollow ugc">notebook</a> and the corresponding <a href="https://github.com/chrisvdweth/nus-cs5246/blob/master/notebooks/src/rnn.py" rel="noopener nofollow ugc">script</a>. Note that it defines the <code>nn.LSTM</code> layer with <code>bidirectional=True</code> and <code>batch_first=True</code>. But conceptually, it’s the same setup/architecture.</p>
<p>Important is also the line</p>
<pre><code class="lang-auto">loss = criterion(outputs.permute(0,2,1), targets)
</code></pre>
<p>more specifically the <code>.permute(0,2,1)</code>. Again, you have to check the <a href="https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html" rel="noopener nofollow ugc">docs</a> for the loss function to see that the expected input shape is <code>(batch_size, target_size, seq_len)</code></p>
<aside class="quote no-group" data-username="yaskovdev" data-post="3" data-topic="202567">
<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/yaskovdev/48/69953_2.png" class="avatar"> yaskovdev:</div>
<blockquote>
<p>So, if I understand correctly, if the output of my <code>nn.LSTM</code> layer has a shape of <code>(BATCH_SIZE, SEQ_LENGTH, HIDDEN_DIM)</code>, then I the arguments of the <code>nn.NLLLoss</code> instance (its input and target) should have the next shape: <code>(BATCH_SIZE * SEQ_LENGTH, HIDDEN_DIM)</code> and <code>(BATCH_SIZE * SEQ_LENGTH)</code>.</p>
</blockquote>
</aside>
<p>I’m fairly certain that this only runs because <code>batch_size=1</code> and wouldn’t otherwise. Of course, have larger batches requires to handle sequences of different lengths – hence, this tutorial ignores this, I would think. In the notebook I link, I use a custom sampler that generates batches that always contain sequences of the same length.</p>

username: yaskovdev
cooked: <p>Thank you for the detailed explanation, <a class="mention" href="/u/vdw">@vdw</a>!</p>
<p>I <a href="https://github.com/jshaipuka/ai-sandbox/blob/master/music-generation/train.py" rel="noopener nofollow ugc">tried</a>:</p>
<ol>
<li>Using <code>batch_size</code> bigger than 1,</li>
<li>Calling my loss function similarly to how you do it in your notebook (<code>criterion(outputs.permute(0,2,1), targets)</code>),</li>
<li>Setting <code>batch_first</code> to its default value and transposing the LSTM input / output correspondingly.</li>
</ol>
<p>Everything worked as expected. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>To be honest, I did not find a good reason for <code>batch_first</code> to be <code>False</code> by default (including this forum): it seems like it only leads to two extra <code>transpose(1, 0)</code> calls before and after the LSTM layer without readability / performance improvements, but that’s probably a separate topic.</p>