Replacing GRU with LSTM in an Encoder-Decoder architecture results in a dimensions mismatch - why?

Lost_Wanderer: <p>Hello PyTorch developers,</p>
<p>I am trying to replace GRU’s with LSTM’s in an Encoder-Decoder architecture and it results in an error. I don’t understand why the error is there. <strong>Could someone shed some light on what is going on here?</strong></p>
<p>Here’s the code which works, along with its outputs. The code pertains to Exercise 5, Chapter 9.7 of the book Dive into Deep Learning:</p>
<pre><code class="lang-auto">#@save
class Seq2SeqEncoder(d2l.Encoder):
    """The RNN encoder for sequence to sequence learning."""
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        # Embedding layer
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers,
                          dropout=dropout)

    def forward(self, X, *args):
        # The output `X` shape: (`batch_size`, `num_steps`, `embed_size`)
        X = self.embedding(X)
        # In RNN models, the first axis corresponds to time steps
        X = X.permute(1, 0, 2)
        # When state is not mentioned, it defaults to zeros
        output, state = self.rnn(X)
        # `output` shape: (`num_steps`, `batch_size`, `num_hiddens`)
        # `state` shape: (`num_layers`, `batch_size`, `num_hiddens`)
        return output, state
</code></pre>
<pre><code class="lang-auto">encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,
                         num_layers=2)
encoder.eval()
X = torch.zeros((4, 7), dtype=torch.long)
output, state = encoder(X)
output.shape
</code></pre>
<p><code>torch.Size([7, 4, 16])</code></p>
<p><code>state.shape</code></p>
<p><code>torch.Size([2, 4, 16])</code></p>
<pre><code class="lang-auto">class Seq2SeqDecoder(d2l.Decoder):
    """The RNN decoder for sequence to sequence learning."""
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqDecoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, num_layers,
                          dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, *args):
        return enc_outputs[1]

    def forward(self, X, state):
        # The output `X` shape: (`num_steps`, `batch_size`, `embed_size`)
        X = self.embedding(X).permute(1, 0, 2)
        print("X.shape")
        print(X.shape)
        # Broadcast `context` so it has the same `num_steps` as `X`
        context = state[-1].repeat(X.shape[0], 1, 1)
        X_and_context = torch.cat((X, context), 2)
        output, state = self.rnn(X_and_context, state)
        output = self.dense(output).permute(1, 0, 2)
        # `output` shape: (`batch_size`, `num_steps`, `vocab_size`)
        # `state` shape: (`num_layers`, `batch_size`, `num_hiddens`)
        return output, state
</code></pre>
<pre><code class="lang-auto">decoder = Seq2SeqDecoder(vocab_size=10, embed_size=8, num_hiddens=16,
                         num_layers=2)
decoder.eval()
state = decoder.init_state(encoder(X))
output, state = decoder(X, state)
output.shape, state.shape
</code></pre>
<pre><code class="lang-auto">X.shape
torch.Size([7, 4, 8])

(torch.Size([4, 7, 10]), torch.Size([2, 4, 16]))
</code></pre>
<p>Now, here is my code. Notice how I just replaced <code>nn.GRU</code> with <code>nn.LSTM</code>:</p>
<pre><code class="lang-auto">#@save
class Seq2SeqEncoder(d2l.Encoder):
    """The RNN encoder for sequence to sequence learning."""
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        # Embedding layer
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.LSTM(embed_size, num_hiddens, num_layers, # not GRU, but LSTM
                          dropout=dropout)

    def forward(self, X, *args):
        # The output `X` shape: (`batch_size`, `num_steps`, `embed_size`)
        X = self.embedding(X)
        # In RNN models, the first axis corresponds to time steps
        X = X.permute(1, 0, 2)
        # When state is not mentioned, it defaults to zeros
        output, state = self.rnn(X)
        # `output` shape: (`num_steps`, `batch_size`, `num_hiddens`)
        # `state` shape: (`num_layers`, `batch_size`, `num_hiddens`)
        return output, state
</code></pre>
<pre><code class="lang-auto">encoder = Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,
                         num_layers=2)
encoder.eval()
X = torch.zeros((4, 7), dtype=torch.long)
output, state = encoder(X)
output.shape
</code></pre>
<p><code>torch.Size([7, 4, 16])</code></p>
<p>Now in the next code cell you can already see some issues:</p>
<p><code>state.shape</code></p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_83249/2608527015.py in &lt;module&gt;
----&gt; 1 state.shape

AttributeError: 'tuple' object has no attribute 'shape'
</code></pre>
<p>Here’s my decoder class:</p>
<pre><code class="lang-auto">class Seq2SeqDecoder(d2l.Decoder):
    """The RNN decoder for sequence to sequence learning."""
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqDecoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.LSTM(embed_size + num_hiddens, num_hiddens, num_layers, # not GRU, but LSTM
                          dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, *args):
        return enc_outputs[1]

    def forward(self, X, state):
        # The output `X` shape: (`num_steps`, `batch_size`, `embed_size`)
        X = self.embedding(X).permute(1, 0, 2)
        print("X.shape:")
        print(X.shape)
        # Broadcast `context` so it has the same `num_steps` as `X`
        context = state[-1].repeat(X.shape[0], 1, 1)
        #print("state[-1].shape:")
        #print(state[-1].shape)
        #context = state[-1]
        X_and_context = torch.cat((X, context), 2)
        output, state = self.rnn(X_and_context, state)
        output = self.dense(output).permute(1, 0, 2)
        # `output` shape: (`batch_size`, `num_steps`, `vocab_size`)
        # `state` shape: (`num_layers`, `batch_size`, `num_hiddens`)
        return output, state
</code></pre>
<pre><code class="lang-auto">decoder = Seq2SeqDecoder(vocab_size=10, embed_size=8, num_hiddens=16,
                         num_layers=2)
decoder.eval()
state = decoder.init_state(encoder(X))
output, state = decoder(X, state)
output.shape, state.shape
</code></pre>
<pre><code class="lang-auto">X.shape:
torch.Size([7, 4, 8])

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_83249/3682229435.py in &lt;module&gt;
      3 decoder.eval()
      4 state = decoder.init_state(encoder(X))
----&gt; 5 output, state = decoder(X, state)
      6 output.shape, state.shape

~/anaconda3/envs/d2l/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/tmp/ipykernel_83249/4080674284.py in forward(self, X, state)
     22         #print(state[-1].shape)
     23         #context = state[-1]
---&gt; 24         X_and_context = torch.cat((X, context), 2)
     25         output, state = self.rnn(X_and_context, state)
     26         output = self.dense(output).permute(1, 0, 2)

RuntimeError: torch.cat(): Sizes of tensors must match except in dimension 2. Got 7 and 14 in dimension 0 (The offending index is 1)
</code></pre>
<p>As you can see, I get a <code>RuntimeError</code>.</p>
<p><strong>What is going on here?</strong> Whenever I was modfying some code by replacing nn.GRU with nn.LSTM, the integration was seamless. I thought that nn.GRU and nn.LSTM differed only in their internal implementation and that their inputs and outptus were the same.</p>
<p>Thank you in advance!</p>

ptrblck: <aside class="quote no-group" data-username="Lost_Wanderer" data-post="1" data-topic="130562">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lost_wanderer/48/41463_2.png" class="avatar"> Lost_Wanderer:</div>
<blockquote>
<p>I thought that nn.GRU and nn.LSTM differed only in their internal implementation and that their inputs and outptus were the same.</p>
</blockquote>
</aside>
<p>No, that’s not the case, as their states are different, which also seems to case the issues.<br>
You can see the expected inputs in the docs for <a href="https://pytorch.org/docs/stable/generated/torch.nn.GRU.html?highlight=gru#torch.nn.GRU"><code>nn.GRU</code></a> and <a href="https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html?highlight=lstm#torch.nn.LSTM"><code>nn.LSTM</code></a>.<br>
E.g. the first error is raised, since the <code>LSTM</code> returns one output tensor and two states.</p>

Lost_Wanderer: <p>Thank you for pointing this out. I have managed to fix the first error by accessing <code>state[0].shape</code> instead of <code>state.shape</code>.</p>
<p>However, my second error still persists. What bothers me is that the dimension of <code>state[0].shape</code> is <code>torch.Size([2, 4, 16])</code>, while the first dimension I need is 7. Since 7 is not divisible by 2, I don’t see a way to make the <code>context</code> variable have 7 as its first dimension.</p>

ptrblck: <p>I don’t know where the required shape of <code>7</code> is coming from, but from the docs:</p>
<blockquote>
<p>h_n: tensor of shape (D∗num_layers,N,H out) containing the final hidden state for each element in the batch.</p>
</blockquote>
<p>Given that you are using <code>num_layers=2</code> and <code>bidirectional=False</code>, <code>dim0</code> is expected to have the shape <code>2</code>.</p>

Lost_Wanderer: <p>This fixed the bug I was having. Thank you.</p>
<p>I noticed that I have a hard time with recursive neural networks - more specifically, with the dimensions of the tensors they expect as input and output. I was reading Dive into Deep Learning, Section 9.7. and it does a great job of explaining the theory behind encoders, decoders and the like. However, I found it lacking in explaining why the dimensions of inputs and outputs are the way they are.</p>
<p><strong>Is there any resource you could recommend to me that focuses on understanding the input and output shapes of various models that are based on GRUs or LSTMs?</strong></p>
<p>Thank you in advance!</p>

ptrblck: <p>I think the <a href="https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html">Seq2Seq tutorial</a> might be a good starter (as well as generally checking the docs to see what shapes are expected).<br>
Besides that <a class="mention" href="/u/vdw">@vdw</a> has shared some excellent posts e.g. <a href="https://discuss.pytorch.org/t/rnn-output-vs-hidden-state-dont-match-up-my-misunderstanding/43280">here</a> and <a href="https://discuss.pytorch.org/t/using-batches-for-seq2seq-models/55205">here</a>, which clarify the structure a bit more.</p>