Using auto_scale_batch_size with variable length batches (sequential data)

[jonathanking]

In my sequence modeling work, I have models that use batches of shape batch_size x sequence_length x D where sequence_length may be different between each batch.

How does the pl.Trainer(auto_scale_batch_size=...) flag work in this case? Am I unable to use this feature? I imagine it may assume that the batches are similar in size except for the batch dimension.

Surely I am not the only person working on modeling tasks with variable-length sequences. If this feature is unavailable in this context, I would love to hear if anyone has written alternatives or modifications that have the same effect.

Thank you for your help!

[rohitgr7]

well, auto_scale_batch_size relies on dataloader, by increasing the batch_size at every iteration. So even if you have variable-length sequences, the batch_size for that dataloader still remains the same for all the batches within that iteration, thus it should work fine.

[jonathanking]

Hmm, I'm not sure I understand.

Are you saying it will test every possible batch (and therefore every possible sequence length) in a DataLoader? That would be a complete epoch and would be too slow.

How many batches are tested with a given batch size?

It sounds like it may work only if you can count on the DataLoader to yield batches with the largest sequence length first. Is that right?

[rohitgr7]

you can control the number of batches per trial using steps_per_trial:
https://github.com/PyTorchLightning/pytorch-lightning/blob/6fcfd8ce91c5e6d5721243c1d8c7606ee3079460/pytorch_lightning/tuner/batch_size_scaling.py#L32-L36

It sounds like it may work only if you can count on the DataLoader to yield batches with the largest sequence length first.

yes, for your use-case.

[jonathanking]

@rohitgr7 Thank you for explaining this to me. I appreciate your help.

In my opinion, this assumption about yielding batches with the largest sequence lengths first is not clear from the documentation. Anyone using variable-length sequential data (like many NLP researchers) would come across this same issue, I imagine.

One problem remains, unfortunately. Let's say I have set up my Datamodule/Dataloader to yield batches with the largest sequence lengths first as you have suggested. This is helpful for determining the largest batch size, but afterward, it would be better for the data to be shuffled.

After running trainer.tune(), could I then turn shuffling on by accessing the Dataloaders directly? I'm curious if you have any advice on how to best approach this.