Implement tensorflow bi-directional code to pytorch

Bishwa_Karki: <pre><code>def encoder():
  enc = keras.Sequential()
  enc.add(Input(shape=(maxlen,), name='Encoder-Input'))
  enc.add(Embedding(num_words, embed_dim,input_length = maxlen, name='Body-Word-Embedding', mask_zero=False))
  enc.add(Bidirectional(LSTM(128, activation='relu', name='Encoder-Last-LSTM')))

  return enc
</code></pre>
<p><a href="https://i.stack.imgur.com/4ciHS.png" rel="noopener nofollow ugc"><img src="https://i.stack.imgur.com/4ciHS.png" alt="enter image description here" width="690" height="332"></a></p>
<p>Here I get output of shape <strong>Batch x 256</strong></p>
<p>In pytorch:</p>
<pre><code>class EncoderRNN(nn.Module):
def __init__(self):
    super(EncoderRNN, self).__init__()
    self.embedding = nn.Embedding(vocab_size, embed_dim)
    self.lstm = nn.LSTM(embed_dim, embed_dim, 2, bidirectional=True, batch_first=False)
    self.relu = nn.ReLU()

def forward(self, input_, hidden, c):
    embedded = self.embedding(input_).transpose(0, 1)
    output, (hn, cn) = self.lstm(embedded, (hidden, c))
    output = self.relu(output)
    return out, hn, cn

def initHidden(self):
    return torch.zeros(2*2, 16, 128, device=device), torch.zeros(2*2, 16, 128, device=device)
</code></pre>
<p>But here I get output of shape: <strong>sequence length x batch x 256</strong> and I want to get <strong>batch x 256</strong></p>
<p>How to implement the exact code of tensorflow to pytorch?</p>

vdw: <p><code>output</code> will contain the last hidden state (“last” w.r.t. the number of layers) for all time steps. You want to use <code>hn</code> which gives you the last hidden states (“last” w.r.t the number of time steps) for all layers. Check out this images below:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/4/748ce2c7f184e70e069326543de260902b9a4e76.png" data-download-href="https://discuss.pytorch.org/uploads/default/748ce2c7f184e70e069326543de260902b9a4e76" title="lstm-output-vs-hidden"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/4/748ce2c7f184e70e069326543de260902b9a4e76.png" alt="lstm-output-vs-hidden" data-base62-sha1="gD3escl8na3s5EFycJH7t67PFJk" width="583" height="500" data-dominant-color="EDEFF2"><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">lstm-output-vs-hidden</span><span class="informations">640×548 21 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>You might also check out the docs to make sure you handle the bidirectional setup correctly.</p>

Bishwa_Karki: <p>Thank you. I am also confused with the shape of last hidden state as it gives output of shape:</p>
<p>[ number of layers x batch x hidden_dim] → if I use 5 layers → [5, 64, 256].</p>
<p>Should it not be [batch, hidden_dim] as I want the latent representation of the input? What does 5 signify when I pass it to decoder?</p>

vdw: <p>You can simply check the <a href="https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html" rel="noopener nofollow ugc">docs</a>: The output shape of <code>hn</code> is <code>(D*num_layers, N, H_out​)</code> where <code>D</code> is the number of directions (1 or 2), <code>N</code> is the batch size, and <code>H_out</code> is your <code>hidden_dim</code>.</p>
<p>Are yo sure you used <code>bidirectional=True</code> to get this result, because with <code>num_layers=5</code>, the shape of <code>hn</code> should be <code>(10, 64, 256)</code> with 10=5*2 reflecting the 5 layers and the 2 directions.</p>

Bishwa_Karki: <p>Yes sorry about the 5, I was talking in terms of general LSTM but yes if we use bi-directional LSTM it’s 10 with 5 number of layers. My question is what does this 10 tell us about when we pass it to decoder? Should the decoder also have to have 5 bi-directional layer?</p>

vdw: <p>It depends on your task. For an autoencoder, for example, you can set up your decoder the same way as your encoder ans simply copy <code>(hn, cn)</code> over.</p>
<p>If you are generating text sequences (e.g., for machine) translation. A bidirectional decoder doesn’t make sense as you generate the the output sequence step by step. In this case you have to see, how to best makes the hidden state from the decoder a valid hidden state of the decoder. The easiest would probably that the decoder LSTM has the same number of layers as the encoder LSTM – but with <code>bidirectional=False</code> and bring the encoder hidden state from <code>(10, 64, 256)</code> to <code>(5, 64, 256)</code> by adding the two directions. It really depends how your decoder looks like.</p>

Bishwa_Karki: <p>thank you so much for your help</p>