How to get the output at the last timestep for batched sequences?

rk2900: <p>I use <code>max_seq_len * batch_size * embed_size</code> for batch input, also with a list of actual lengthes for each sequence to GRU/LSTM,<br>
then I get the outputs and last hidden vector with size of <code>max_seq_len * batch_size * hidden_size</code> and <code>layer_num * batch_size * hidden_size</code>.<br>
How can I get the actual output vector at the last (actual) time step for each sequence? whose size should be <code>batch_size * hidden_size</code>.</p>
<p>The test code for my question:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.autograd import Variable
from torch import optim
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
import numpy as np

max_seq_len = 3
batch_size = 4
layer_num = 2

input_size = 10 # 0 - 9
emb_size = 8
hidden_size = 10

# test data:
a = [[1,2,3], [4,5,0], [6,0,0], [7,8,9]]
lens = [3,2,1,3]

# sort the input batch data by reversed actual length for pad_pack operation
pairs = sorted( zip(a, lens), key=lambda p: p[1], reverse=True)
(a, lens) = zip(*pairs)

# actual length
lens = np.array(lens)
# lens = torch.LongTensor(lens)

va = Variable(torch.LongTensor(a))
vlens = Variable(torch.LongTensor(lens))

print(va, vlens)

embedding = nn.Embedding(input_size, emb_size)
gru = nn.GRU(emb_size, hidden_size, layer_num)

inputs = va.transpose(0, 1); print("inputs size: ", inputs.size()) # max_seq_len * batch_size
inputs = embedding(inputs); print("embedded size: ", embedded.size()) # max_seq_len * batch_size * emb_size

inputs = torch.nn.utils.rnn.pack_padded_sequence(inputs, lens)

h0 = Variable(torch.randn(layer_num, 4, 10))
outputs, hn = gru(inputs, h0)

# print("after packed:")
# print("outputs.size: ", outputs.size())
# print("hn size: ", hn.size())

tmp = torch.nn.utils.rnn.pad_packed_sequence(outputs)

outputs, output_lengths = tmp
# outputs: (max_seq_len * batch_size, output_size)

print("after padded:")
print("outputs.size: ", outputs.size())
print("hn size: ", hn.size())

# outputs.index_select(torch.LongTensor(lens-1))
# print("lens: ", lens-1)

# idxs = Variable(torch.LongTensor(lens-1))
# outputs = outputs.index_select(0, idxs)

# print("selected outputs size: ", outputs.size())
# outputs.gather(0, idxs)
</code></pre>

rk2900: <p>My direct question is</p>
<p>How to implement <code>[ outputs[lens[i]][i] for i in range(len(lens)) ]</code> ?<br>
which is to get the last output vector at the last timestep for each sequence.</p>
<p>Thank you.</p>

rk2900: <p>My own solution is</p>
<pre><code class="lang-auto">masks = (vlens-1).unsqueeze(0).unsqueeze(2).expand(max_seq_len, outputs.size(1), outputs.size(2))
output = outputs.gather(0, masks)[0]
</code></pre>

eugeneware: <p><a class="mention" href="/u/rk2900">@rk2900</a></p>
<p>Thanks for posting this. I think I’ve found an easier way. For GRUs, the last hidden state is equivalent to the last output state I believe.</p>
<p>So you should be able to do:</p>
<pre><code class="lang-auto">outputs, hn = gru(inputs, h0)
print(hn[-1])
</code></pre>
<p>For the LSTM the equivalent code would be:</p>
<pre><code class="lang-auto">outputs, (hn, cn) = lstm(inputs, h0)
print(hn[-1])
</code></pre>
<p>I used your code to verify this.</p>
<p>And you can more compactly express your code by using <code>.view()</code> to add the unit axes:</p>
<pre><code class="lang-auto">masks = (vlens-1).view(1, -1, 1).expand(max_seq_len, outputs.size(1), outputs.size(2))
output = outputs.gather(0, masks)[0]
</code></pre>

bsantraigi: <p>Why are we expanding in the 0th dimension also? Wouldn’t the following return the same thing? If they are same, is the following faster? (because of less number of gather calls)</p>
<pre><code class="lang-auto">masks = (vlens-1).unsqueeze(0).unsqueeze(2).expand(1, outputs.size(1), outputs.size(2))
output = outputs.gather(0, masks)
</code></pre>

mrusso: <p>I have a question related to this post so I thought I would follow up here rather than creating a new post. I’ve currently been doing the following to extract the output at the last (non-padded) timestep for an LSTM I’m using (note: <code>batch_first</code> is <code>True</code> here):</p>
<pre><code class="lang-auto">def encoder(self, input: torch.Tensor, lengths: torch.Tensor)
    # pack our zero-padded sequences
    packed_seq = nn.utils.rnn.pack_padded_sequence(input, lengths, batch_first=True)

    # pass input through lstm
    packed_output, _ = self.lstm(packed_seq)

    # unpack the output which is (batch x seq_len x num_directions * hidden_size)
    total_length = input.size(1)
    output, _ = torch.nn.utils.rnn.pad_packed_sequence(packed_output, batch_first=True, total_length=total_length)

    # grab the final output for each sequence in the batch
    seq_len_indices = [length - 1 for length in lengths]
    batch_indices = [i for i in range(self.batch_size)]
    final_output = output[batch_indices, seq_len_indices, :]

    return final_output
</code></pre>
<p>My <code>final_output</code> has the right dimensions (<code>batch_size x hidden_size</code>) and seems to get output from the right timesteps, but I’m wondering if I need to be using the torch operations people used above^ (e.g. the <code>gather</code> function with <code>masks</code>) in order for autograd to backpropogate the loss correctly? Will slicing the <code>output</code> tensor using regular lists mess up the backprop?</p>

btimar: <p>Slicing via lists doesn’t mess up any gradient computations. Explicitly:</p>
<pre><code class="lang-python">import torch
from torch.nn.utils.rnn import pad_sequence

h1 = torch.tensor([1., 1, 2, 3], requires_grad=True)
h2 = torch.tensor([5., 5], requires_grad=True)
h = pad_sequence([h1, h2])
lengths = torch.tensor([len(h1), len(h2)]).long()
batch_size = len(lengths)

# list slicing
final_states_list = h[[l-1 for l in lengths], [i for i in range(batch_size)]]

# gather via mask
mask = (lengths-1).view(1, -1).expand(h.size(0), h.size(1))
final_states_mask = h.gather(0, mask)[0]

#check grads both ways
final_states_list.sum().backward(retain_graph=True)
print(h1.grad, h2.grad)
# tensor([0., 0., 0., 1.]) tensor([0., 1.])
h1.grad.zero_()
h2.grad.zero_()
final_states_mask.sum().backward()
print(h1.grad, h2.grad)
# tensor([0., 0., 0., 1.]) tensor([0., 1.])
</code></pre>
<p>I’m not sure what the advantage of using gather() here is, since all rows but the first of the gather()'d tensor are thrown away. The list slicing seems cleaner.</p>

mrusso: <p>This is a great response, thank you for the example it was very informative!</p>