Get each sequence's last item from packed sequence

username: BramVanroy
cooked: <p><a href="https://stackoverflow.com/questions/55399115/get-each-sequences-last-item-from-packed-sequence" rel="nofollow noopener">[x-post from Stack Overflow]</a></p>
<p>I am trying to put a packed and padded sequence through a GRU, and retrieve the output of the last item of each sequence. Of course I don’t mean the <code>-1</code> item, but the actual last, not-padded item. We know the lengths of the sequences in advance, so it should be as easy as to extract for each sequence the <code>length-1</code> item.</p>
<p>I tried the following</p>
<pre><code class="lang-python">import torch
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Data
input = torch.Tensor([[[0., 0., 0.],
                       [1., 0., 1.],
                       [1., 1., 0.],
                       [1., 0., 1.],
                       [1., 0., 1.],
                       [1., 1., 0.]],
                      
                      [[1., 1., 0.],
                       [0., 1., 0.],
                       [0., 0., 0.],
                       [0., 1., 0.],
                       [0., 0., 0.],
                       [0., 0., 0.]],

                      [[0., 0., 0.],
                       [1., 0., 0.],
                       [1., 1., 1.],
                       [0., 0., 0.],
                       [0., 0., 0.],
                       [0., 0., 0.]],

                      [[1., 1., 0.],
                       [0., 0., 0.],
                       [0., 0., 0.],
                       [0., 0., 0.],
                       [0., 0., 0.],
                       [0., 0., 0.]]])

lengths = [6, 4, 3, 1]
p = pack_padded_sequence(input, lengths, batch_first=True)

# Forward
gru = torch.nn.GRU(3, 12, batch_first=True)
packed_output, gru_h = gru(p)

# Unpack
output, input_sizes = pad_packed_sequence(packed_output, batch_first=True)

last_seq_idxs = torch.LongTensor([x-1 for x in input_sizes])

last_seq_items = torch.index_select(output, 1, last_seq_idxs) 

print(last_seq_items.size())
# torch.Size([4, 4, 12])
</code></pre>
<p>But the shape is not what I expect. I had expected to get <code>4x12</code>, i.e. <code>last item of each individual sequence x hidden</code>.`</p>
<p>I could loop through the whole thing, and build a new tensor containing the items I need, but I was hoping for a built-in approach that took advantage of some smart math. I fear that manually looping and building, will result in very poor performance.</p>

username: TrentBrick
cooked: <p>No idea if this is the most efficient method but I think it gets the job done.</p>
<p>Change your line:</p>
<pre><code class="lang-auto">last_seq_items = torch.index_select(output, 1, last_seq_idxs) 
</code></pre>
<p>to this:</p>
<pre><code class="lang-auto">last_seq_items = output[range(output.shape[0]), last_seq_idxs, :]
</code></pre>
<p>By using <code>range()</code>, we are able to give it a different index to take at each 0th dimension.</p>

username: BramVanroy
cooked: <p><a class="mention" href="/u/trentbrick">@TrentBrick</a> Thank you for that. That makes sense, I think.</p>
<p>Personally, I came up with this:</p>
<pre><code class="lang-python">output, input_sizes = pad_packed_sequence(packed_output, batch_first=True)
# One per sequence, with its last actual node extracted, and unsqueezed
last_seq = [output[e, i-1, :].unsqueeze(0) for e, i in enumerate(input_sizes)]
# Merge them together
last_seq = torch.cat(last_seq, dim=0) 
</code></pre>

username: TrentBrick
cooked: <p><a class="mention" href="/u/bramvanroy">@BramVanroy</a> I think mine might be faster because it is for loop free! And glad I could help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=7" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: BramVanroy
cooked: <p><a href="https://stackoverflow.com/a/55402149/1150683" rel="nofollow noopener">Someone</a> suggested using torch’s <code>arange</code>. Not sure whether this would be faster than Python’s native range, but I am posting it for completeness sake.</p>

username: TrentBrick
cooked: <p>That is probably right!</p>

username: aRI0U
cooked: <p>I found a way to do this which is about 100x faster and is also better in terms of memory usage. The main idea is to directly extract the indices from the <code>PackedSequence</code> itself instead of padding it.</p>
<p>Assuming you have a <code>PackedSequence</code> object named <code>packed</code> and containing sequences of respective length <code>lengths</code>, you can extract the last item of each sequence this way:</p>
<pre><code class="lang-auto">sum_batch_sizes = torch.cat((
    torch.zeros(2, dtype=torch.int64),
    torch.cumsum(packed.batch_sizes, 0)
))
sorted_lengths = lengths[packed.sorted_indices]
last_seq_idxs = sum_batch_sizes[sorted_lengths] + torch.arange(lengths.size(0))
last_seq_items = packed.data[last_seq_idxs]
last_seq_items = last_seq_items[packed.unsorted_indices]
</code></pre>
<p>For the most skeptical of you, here is a sketch of the proof:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/e/0e54f84b0ce4d19dfd2e8e98045ece0e52bc61c9.png" data-download-href="https://discuss.pytorch.org/uploads/default/0e54f84b0ce4d19dfd2e8e98045ece0e52bc61c9" title="proof"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/e/0e54f84b0ce4d19dfd2e8e98045ece0e52bc61c9.png" alt="proof" data-base62-sha1="22MJmPPctql17UR35Vxr9QilgAp" width="690" height="433" data-dominant-color="F9F9F9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">proof</span><span class="informations">1054×662 24.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>With the notations introduced above:</p>
<ul>
<li><code>packed.data</code> corresponds to <em>X</em></li>
<li><code>packed.batch_sizes</code> corresponds to [<em>B_0, B_1, …, B_l0-1</em>]</li>
<li><code>packed.sorted_indices</code> and <code>packed.unsorted_indices</code> correspond to the permutation so that sequences are sorted by decreasing length</li>
</ul>
<p>Hence the code above.</p>
<p>Since there is no need to pad the packed sequence to extract the last items with that method, it is faster than the others proposed. It is also more convenient if you have limited memory usage since you don’t have to store all the zeros from the padded sequences.</p>
<p>Here is a test and benchmark for 10000 sequences with random lengths between 1 and 100:</p>
<pre><code class="lang-auto">Using CPU
Method 1 # TrentBrick Mar 28, 2019
Error: 0.0
10 loops, best of 3: 97.1 ms per loop

Method 2 # BramVanroy Mar 28, 2019
Error: 0.0
10 loops, best of 3: 171 ms per loop

Method 3 # mine
Error: 0.0
1000 loops, best of 3: 820 µs per loop

Using GPU
Method 1
Error: 0.0
10 loops, best of 3: 69 ms per loop

Method 2
Error: 0.0
10 loops, best of 3: 168 ms per loop

Method 3
Error: 0.0
1000 loops, best of 3: 430 µs per loop
</code></pre>
<p>Here is the whole script:</p>
<pre><code class="lang-auto">import torch

from torch.nn.utils.rnn import pack_sequence, pad_packed_sequence

lengths = torch.randint(1, 100, (10000,))

sequences = [torch.randn(i, 2) for i in lengths]

ground_truth = torch.stack([seq[-1] for seq in sequences])

packed = pack_sequence(sequences, enforce_sorted=False)

def method1(packed): # TrentBrick Mar 28, 2019
    output, input_sizes = pad_packed_sequence(packed, batch_first=True)
    last_seq_idxs = torch.LongTensor([x-1 for x in input_sizes])
    last_seq_items = output[range(output.shape[0]), last_seq_idxs, :]
    return last_seq_items

def method2(packed): # BramVanroy Mar 28, 2019
    output, input_sizes = pad_packed_sequence(packed, batch_first=True)
    last_seq_items = [output[e, i-1, :].unsqueeze(0) for e, i in enumerate(input_sizes)]
    last_seq_items = torch.cat(last_seq_items, dim=0)
    return last_seq_items

def method3(packed): # mine
    sum_batch_sizes = torch.cat((
        torch.zeros(2, dtype=torch.int64),
        torch.cumsum(packed.batch_sizes, 0)
    ))
    sorted_lengths = lengths[packed.sorted_indices]
    last_seq_idxs = sum_batch_sizes[sorted_lengths] + torch.arange(lengths.size(0))
    last_seq_items = packed.data[last_seq_idxs]
    last_seq_items = last_seq_items[packed.unsorted_indices]
    return last_seq_items

print('Using CPU')
ground_truth = ground_truth.cpu()
packed = packed.cpu()
lengths = lengths.cpu()
for i, method in enumerate([method1, method2, method3]):
    print('Method', i+1)
    print('Error:', torch.norm(ground_truth - method(packed)).item())
    %timeit method(packed)
    print()

print('Using GPU')
ground_truth = ground_truth.cuda()
packed = packed.cuda()
# lengths = lengths.cuda()
for i, method in enumerate([method1, method2, method3]):
    print('Method', i+1)
    print('Error:', torch.norm(ground_truth - method(packed)).item())
    %timeit method(packed)
    print()
</code></pre>