Pad PackedSequences to original batch length

[ajfisch]

The current flow for handling variable length sequences in RNNs is:

packed_input = torch.nn.utils.rnn.pack_padded_sequence(input, lengths)
packed_output = rnn(packed)[0]
output = torch.nn.utils.rnn.pad_packed_sequence(packed_output)

The output size of the Variable returned by pad_packed_sequence is determined by the max length in lengths: https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/rnn.py#L106.

However, this only works in recovering the original size of the input if the max length sequence has no padding (max length == length dim of batched input). For normal, sensible batching this makes sense and should be true.

But if a model is using, say, DataParallel, the batch might be split such that there is extra padding. And the output size from the RNN will be truncated (which might break other things down-stream).

To fix this potentially unexpected behavior, I propose two possible simple patches.

1. A max_batch_size field is calculated from the original input and added to the PackedSequence namedtuple: https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/rnn.py#L6 and used instead of https://github.com/pytorch/pytorch/blob/master/torch/nn/utils/rnn.py#L106

2. An optional max_batch_size parameter is added to pad_packed_sequence which would be used as an override.

I prefer 1), but I suppose 2) has the advantage of being fully backwards compatible (even though it is discouraged to be directly creating or meddling with PackedSequence tuples, it is possible).

cc @cpuhrsch

[jekbradbury]

I agree with your preferred solution even though I think my team is relying on this behavior (which also tripped us up) in our code at the moment

[yikang-li]

Hi, I encounter the same problem when using DataParallel. Do you have any solution to the problem?

[manasRK]

This issue is tripping up our pipeline as well. For now, we are concat-ing dummy tensors and masking to take care of this.

However, is anyone working on this now?

[soumith]

@manasRK no one's working on it, you can go for it.

[yikang-li]

Our solution is to use original pad_packed_sequence and then add further padding to the original lengths.

[manasRK]

For people who stumble upon this thread in the future, this is how we are handling it currently (for now);

First, the usual pack_padded_sequence and pad_packed_sequence for handling variable length sequences;

seq_len, bsz, n_dims = feats.size()
packed_input = pack_padded_sequence(feats, lengths, batch_first=False)
packed_output, self.hidden = self.lstm(packed_input, self.hidden)
# lstm_out --> seqlen X bsz X hidden_dim
lstm_out, output_lengths = pad_packed_sequence(packed_output, batch_first = False)

Then the hack is implemented as the output size of the Variable returned by pad_packed_sequence is determined by the max length in output_lengths, not seqlen in batch. Also, you may have to hardcode MAXLEN in sequence/loss masking procedures;

if lstm_out.size(0)<seq_len:
    dummy_tensor = autograd.Variable(torch.zeros(seq_len-lstm_out.size(0), bsz, self.hidden_dim))
    lstm_out = torch.cat([lstm_out, dummy_tensor], 0)

Our accuracy metrics have remained stable and predictions in line with our expectations, so I think this hack works well.

[cpuhrsch]

Since we have a solution I'm using this issue to track work that could be replaced by NestedTensors and therefore am assigning myself.