Pytorch equivalent to keras.layers.LSTM(return_sequences=False)

rexdouglass: <p>Keras’s LSTM layer includes a single flag to flatten the output into 1xN-hidden dimensions.<br>
<a href="https://keras.io/layers/recurrent/" class="onebox" target="_blank" rel="nofollow noopener">https://keras.io/layers/recurrent/</a><br>
" * <strong>return_sequences</strong> : Boolean. Whether to return the last output in the output sequence, or the full sequence."</p>
<p>This allows you to process a sequence, convert it to a single embedding, and then pass that to something like a classifier. I don’t want to go token by token. I want to take a sequence of arbitrary length and flatten it to a fixed embedding.</p>
<p>I know there are guides and explanations for how to do this in Pytorch but there are many mistakes and contradictions across different sources. There is a lot of quick speculation, and I’d like a definitive answer if it exists in documentation or example code.</p>
<p>The best diagram I’ve found is here<br>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/48302810/whats-the-difference-between-hidden-and-output-in-pytorch-lstm" target="_blank" rel="nofollow noopener">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/1119340/nathaniel" target="_blank" rel="nofollow noopener">
    <img alt="Nathaniel" src="https://i.stack.imgur.com/ts6WE.jpg?s=128&amp;g=1" class="thumbnail onebox-avatar" width="128" height="128">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/48302810/whats-the-difference-between-hidden-and-output-in-pytorch-lstm" target="_blank" rel="nofollow noopener">What's the difference between "hidden" and "output" in PyTorch LSTM?</a>
</h4>

<div class="tags">
  <strong>pytorch</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/1119340/nathaniel" target="_blank" rel="nofollow noopener">
    Nathaniel
  </a>
  on <a href="https://stackoverflow.com/questions/48302810/whats-the-difference-between-hidden-and-output-in-pytorch-lstm" target="_blank" rel="nofollow noopener">01:54PM - 17 Jan 18 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Even for people writing tutorials there is confusion about what to do<br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/yunjey/pytorch-tutorial/issues/149" target="_blank" rel="nofollow noopener">github.com/yunjey/pytorch-tutorial</a>
  </header>
  <article class="onebox-body">
    <a href="https://github.com/hwijeen" rel="nofollow noopener">
<img src="https://avatars1.githubusercontent.com/u/29157715?v=2&amp;s=96" class="thumbnail onebox-avatar" width="96" height="96">
</a>

<h4><a href="https://github.com/yunjey/pytorch-tutorial/issues/149" target="_blank" rel="nofollow noopener">Issue: output of bidirectional LSTM</a></h4>

<div class="date" style="margin-top:10px;">
	<div class="user" style="margin-top:10px;">
	opened by <a href="https://github.com/hwijeen" target="_blank" rel="nofollow noopener">hwijeen</a>
	on <a href="https://github.com/yunjey/pytorch-tutorial/issues/149" target="_blank" rel="nofollow noopener">2018-12-06</a>
	</div>
	<div class="user">
	</div>
</div>

<pre class="content" style="white-space: pre-wrap;">First of all, thanks for your great tutorial on pytorch! It's a great tip for beginners.
I have a question about the...</pre>

<div class="labels">
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

Shivam_Chandhok: <p>Can you please explain how to do this in pytorch</p>

Humam_Helfawi: <p>Hi ,</p>
<p>Did you find a solution? Thanks!</p>

vdw: <p>The output of an <code>nn.LSTM</code> is <code>output, (h_n, c_n)</code> with the following shapes</p>
<ul>
<li>
<code>output.shape</code>: (seq_len, batch, num_directions * hidden_size)</li>
<li>
<code>h_n.shape</code>:  (num_layers * num_directions, batch, hidden_size)</li>
<li>
<code>c_n.shape</code>:  (num_layers * num_directions, batch, hidden_size)</li>
</ul>
<p>That means that <code>output</code> contains the full sequence as indicated by <code>seq_len</code>.</p>

Humam_Helfawi: <p>Thanks <a class="mention" href="/u/vdw">@vdw</a>! Did you mean that it is enough to take the last part of the sequence to match Keras  output? Did I understood correctly?</p>

vdw: <p>I’m not that familiar with Keras. Maybe <a href="https://discuss.pytorch.org/t/clarification-regarding-the-return-of-nn-gru/47363">this post</a> will help?</p>

Kirtikumar_Pandya: <p>Okay. I found the answer.</p>
<p><code>rnn = nn.LSTM(features_in=10, features_out=20, num_layers=1, batch_first=True)</code><br>
is similar to <code>lstm = tf.keras.layers.LSTM(features_out=20)</code></p>
<p><strong>Note:</strong>  keras does not provide option for how many LSTM layers you want stack therefore I put 1 for <code>num_layers.</code> In keras you don’t have to provide <code>features_in.</code></p>
<pre><code class="lang-auto">rnn = nn.LSTM(10, 20, 1, batch_first=True)

#  [batch, sequence length, features]
input = torch.randn(5, 3, 10)

h0 = torch.randn(1, 5, 20)
c0 = torch.randn(1, 5, 20)
output, (hn, cn) = rnn(input, (h0, c0))

# output shape: [batch, sequence length, out features]
output.shape
</code></pre>
<blockquote>
<p>torch.Size([5, 3, 20])</p>
</blockquote>
<p>Now, for keras:</p>
<pre><code class="lang-auto"># [batch, timesteps, feature]
inputs = tf.random.normal([5, 3, 10])
lstm = tf.keras.layers.LSTM(20, return_sequences=True)
output = lstm(inputs)

# output shape: [batch, sequence length, out features]
output.shape
</code></pre>
<blockquote>
<p>TensorShape([5, 3, 20])</p>
</blockquote>