Best practices for batching data for stateful LSTM and text generation

codingGoose: <p>I’m building an LSTM network for text generation and will train it using the first chapter of a book. My understanding is that this is a good use case for a stateful LSTM since this represents document-level prediction. However, I’m having issues understanding what the best structure for batches is and how to create a custom <code>DataLoader</code> for this purpose.</p>
<p>One approach I’ve though of is to create batches that have a continuous part of the chapter’s text and are shaped <code>(batch_size, timesteps)</code> (timesteps = words). That is, if you merged the batches by their columns, you get a part of the chapter’s text. Here’s an example of what the batch inputs would look like:</p>
<p><strong>Batch <span class="hashtag-raw">#1</span></strong><br>
--------------------------------Timesteps --------------------<br>
<strong>batch element 1|</strong> word_1 | word_2 | word_3 |<br>
<strong>batch element 2|</strong> word_2 | word_3 | word_4 |<br>
<strong>batch element 3|</strong> word_3 | word_4 | word_5 |<br>
…<br>
<strong>Batch <span class="hashtag-raw">#2</span></strong><br>
--------------------------------Timesteps --------------------<br>
<strong>batch element 1|</strong> word_4 | word_5 | word_6 |<br>
<strong>batch element 2|</strong> word_5 | word_6 | word_7 |<br>
<strong>batch element 3|</strong> word_6 | word_7 | word_8 |</p>
<p>The target batch would have the timesteps shifted by +1 so the model predicts the next word for the corresponding timestep in the input batch. Note each row in a batch starts +1 timestep from the above row.</p>
<p>I believe this approach would allow to pass the LSTM states between batch iterations, since the <em>i-th</em> row of a given batch has the text that precedes the text of the next batch’s <em>i-th</em> row. The states resulting form the LSTM layer have shape <code>(num_layers, batch_size, hidden_out)</code>, so there is correspondence across the batch dimension in the inputs and the LSTM states.</p>
<p>I tried to create custom <code>Dataset</code> and <code>DataLoader</code> classes for this purpose but have gotten weird results. Specifically, my questions are:</p>
<ol>
<li>Is the mentioned batch structure a good approach for creating batches for training a text generation LSTM?</li>
<li>Are custom classes of <code>DataLoader</code> a good approach here or is it better to create a tensor with the input info and iterate over it with indexes in the training function?</li>
</ol>
<p>Thanks!</p>

vdw: <p>Yes, this organization of batches should be the correct way to go if you want to “preserve” the hidden state between batches. You probably want to <code>detach()</code> the hidden state though to avoid your computational graph exploding.</p>
<p>You can probably implement this using a <code>Dataset</code> and a <code>Sampler</code>. Here is a <a href="https://discuss.pytorch.org/t/tensorflow-esque-bucket-by-sequence-length/41284/13">basic example</a> that uses this approach to create batches where all sequences in a batch are guaranteed to have the same length. As you can see, the relevant logic is in the <code>Sampler</code>  class.</p>

codingGoose: <p>Thanks <a class="mention" href="/u/vdw">@vdw</a>! I’ll keep in mind your example with <code>Sampler</code> for future implementations. I actually managed to implement this in a different way with an approach I found on a book: They took the encoded text (list of embedding indexes representing the text) and divide it into chunks, with each chunk sliding +1 word with respect to the previous one. The chunks are sized <code>seq_len</code> +1, so that each chunk contains the input sequence for the model and the target sequence (input: <code>[0:-1]</code>, target: <code>[1:]</code>). A <code>Dataset</code> custom class retrieves a chunk from the chunks list by index and returns a tuple with the input and target for the model. This custom <code>Dataset</code> is passed to <code>DataLoader</code> with <code>drop_last=True</code>, generating mini-batches that are contiguous as I described above.</p>