`unpack_sequence` without padding

aRI0U: <p>Hi,<br>
I have operations that I need to do in parallel on the GPU for batches of sequences.<br>
The issue is that the sequences have very different sizes and I have limited GPU resources so I definitely cannot use padding to process my sequences per batch.<br>
I found that <code>PackedSequence</code> could be helpful since data contained in these objects are smartly stored without padding, however at the end I need to unpack the object to get a list of tensors. And <code>unpack_sequence</code> internally uses <code>pad_packed_sequence</code>, whereas the reason why I used PackedSequence is to avoid padding…<br>
Does anyone know if there is a smart way to unpack sequences without naive padding?</p>
<p>Thanks a lot!</p>

aRI0U: <p>Managed to find a solution based on the solution suggested in <a href="https://discuss.pytorch.org/t/get-each-sequences-last-item-from-packed-sequence/41118/7">this post</a>. I post it here for reference.</p>
<p>Here the idea is quite similar, we directly retrieve the elements from the <code>PackedSequence</code> instead of padding it and unpad the results, which avoids putting useless tons of zeros in memory. Assuming a set of sequences of respective lengths <code>lengths</code> packed in a <code>PackedSequence</code> object called <code>packed</code>, the sequences can be retrieved with following function:</p>
<pre><code class="lang-auto">def unpack_sequence(packed: PackedSequence, lengths: torch.Tensor) -&gt; List[torch.Tensor]:
    sum_batch_sizes = pad(packed.batch_sizes.to(packed.data.device), (1, 0)).cumsum(dim=0)
    sequences = []
    for seq_idx, seq_length in zip(packed.unsorted_indices, lengths):
        indices = sum_batch_sizes[:seq_length] + seq_idx
        sequences.append(packed.data[indices])
    return sequences
</code></pre>
<p>This function induces almost no additional GPU memory usage, and is slightly faster than the proposed one for big/high-dimensional sequences (which is usually what you deal with when you have memory issues)</p>
<p>Here is a basic script to test which one is faster for you:</p>
<pre><code class="lang-auto">from typing import List

import torch
from torch.nn.functional import pad
from torch.nn.utils.rnn import PackedSequence


def unpack_sequence(packed: PackedSequence, lengths: torch.Tensor) -&gt; List[torch.Tensor]:
    sum_batch_sizes = pad(packed.batch_sizes.to(packed.data.device), (1, 0)).cumsum(dim=0)
    sequences = []
    for seq_idx, seq_length in zip(packed.unsorted_indices, lengths):
        indices = sum_batch_sizes[:seq_length] + seq_idx
        sequences.append(packed.data[indices])
    return sequences


if __name__ == "__main__":
    import time

    from torch.nn.utils.rnn import pack_sequence
    from torch.nn.utils.rnn import unpack_sequence as naive_unpack_sequence

    device = torch.device("cuda:0")

    num_sequences = 32
    lengths = torch.randint(1, 4096, (num_sequences,))
    seqs = [torch.randint(10, (length, 2048), device=device) for length in lengths]

    packed = pack_sequence(seqs, enforce_sorted=False)

    t0 = time.time()
    unpacked_mine = unpack_sequence(packed, lengths)
    t1 = time.time()
    print("mine:", t1 - t0)

    t0 = time.time()
    unpacked_naive = naive_unpack_sequence(packed)
    t1 = time.time()
    print("naive:", t1 - t0)

    # print(unpacked_naive)
    # print(unpacked_mine)

    for a, b in zip(unpacked_naive, unpacked_mine):
        torch.testing.assert_close(a, b)
    print("OK.")
</code></pre>