[Q] Can someone explain the logic behind tbptt_split_batch's splitting dimension?

[garrett361]

I am confused about the code in tbptt_split_batch referenced at the bottom of this post, specifically why it is apparently assumed that the x in for x in batch would have their time-dimension be at 1.

I would have thought that tbptt_split_batch would be designed with torch's RNN, GRU, or LSTM in mind (with batch_first=True), in which case I would expect batch to be of shape (batch_size, sequence_length, input_size). But if this were the case, then the time_dims = [len(x[0]) for x in batch] would mistakenly take input_size to be the size of the time dimension and the rest oftbptt_split_batch would split along this dimension, rather than the sequence_length dimension, no?

Came across the above when attempting to use tbptt_split_batch for a custom dataset. Needed to overwrite with the splitting over the dimension which I expected, as outline above, and it seems to work correctly.

I feel like I'm badly misunderstanding something.

https://github.com/PyTorchLightning/pytorch-lightning/blob/c9bc10ce8473a2249ffa4e00972c0c3c1d2641c4/pytorch_lightning/core/lightning.py#L1720-L1739

Edit addition for clarity: in the above code I assume that batch is a (b, t, d) shaped input tensor with these three numbers being the batch size, sequence length, and input dimension, respectively.

time_dims = [len(x[0]) for x in batch if isinstance(x, (torch.Tensor, collections.Sequence))]

Then, it would seem that len(x[0]) would give d, rather than the expected t, for each x. I would have thought that the above should read

time_dims = [x.shape[0] for x in batch if isinstance(x, (torch.Tensor, collections.Sequence))]

which is what I have in my own tbptt_split_batch methods. I guess different assumptions are being made about the shape of the input tensor?