Batching and outputting multiple objects with IterableDataset

lucf: <p>Hi All,</p>
<p>I have a very specific application where I use grammar rules to generate random sequences of tokens (up to a maximum length), and train a VAE with these tokens, as in <a href="https://arxiv.org/pdf/1703.01925.pdf" rel="nofollow noopener">https://arxiv.org/pdf/1703.01925.pdf</a>. I generate sequences on the fly in an IterableDataset, which converts the token sequences to collections of one-hot vectors (one per sequence timestep). Note that the VAE is defined as a Pytorch-Lightning model. Here’s the outline of my custom dataset class:</p>
<pre><code class="lang-auto">class SentenceGenerator(IterableDataset):
    def __init__(self, grammar, min_sample_len, max_sample_len, batch_size=256, seed=0):
        super().__init__()
        ...

    def generate_sentence(self):
        # generates a string of tokens ('sentence') using the rules encoded in self.grammar
        ...
        return ''.join(sent)

    def generate_one_hots(self):  
        # converts 'sentences' to sequences of one-hot vectors
        self.sents = [self.generate_sentence() for _ in range(self.batch_sz)]
        out = make_one_hot(self.grammar, self.tokenizer, self.prod_map, self.sents, max_len=self.max_len,
                           n_chars=self.n_chars)
        return out.transpose(2, 1)  # (batch_size, vocab_size, max_length)

    def __iter__(self):
        return iter(self.generate_one_hots()) 
</code></pre>
<ol>
<li>
<p><strong>Batching:</strong> For some reason, generating a single matrix of one-hot vectors and letting the DataLoader batch them didn’t work; the batches were always of size one, plus I was having other downstream issues with Pytorch-Lightning. Therefore I resorted to handling batching directly within <code>SentenceGenerator</code>, as you can see above. In the DataLoader, I then have to specify the same batch size as in the Dataset for batches to be generated. It is a bit hacky and causes some headaches again downstream in terms of understanding what an epoch is, when to step the LR scheduler, when to log a result, etc. Is there a more elegant way of achieving the same result?</p>
</li>
<li>
<p><strong>Multiple outputs from the dataset:</strong> In order to help the VAE train, I would like to explicitly pass it the sentence lengths. This cannot be directly derived from the matrix of one-hot vectors because each token can take a variable number of timesteps to produce. Therefore, I need to output this value directly from the IterableDataset like so:</p>
</li>
</ol>
<pre><code class="lang-auto">    def generate_one_hots(self):
        ...  # same as above except for `return` statement
        return out.transpose(2, 1), self.sent_lengths  # now returning 2 tensors

    def _iter__(self):
         return iter(self.generate_one_hots())
</code></pre>
<p>… and then upack each batch in <code>vae.forward()</code>  as <code>x, n = batch</code> where <code>n</code> contains the lengths of the ‘sentences’ represented by the tensor of one-hot vectors <code>x</code>. However, in this implementation, only the first object returned by <code>generate_one_hots()</code> is included in batch and I get a “too few values to unpack” error. Would you have any suggestions? Ideally one that solves both issues at once!</p>
<p>Thanks a million in advance.</p>