RNNCell uses much more memory than RNN

Matthew_Keaton: <p>Hi everyone! I am working with somewhat complicated RNN architectures that receive inputs from multiple sources in such a way that requires me to process each RNN layer separately and in a sequential fashion. Because of this, I am utilizing the <code>RNNCell</code> block. However, I’ve noticed that when training these networks, the memory utilization far exceeds those of RNNs with the same architectures. As a simple experiment, I devised a 1-layer RNN using the classic <code>RNN()</code> and compared it to an <code>RNNCell()</code>. This is what my sample code looks like:</p>
<p>For <code>RNN</code> -</p>
<pre><code class="lang-auto">class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(RNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, nonlinearity = 'relu', batch_first = True)
        self.fc = nn.Linear(hidden_size, num_classes)


    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)
        print(h0.shape)
        output, hidden = self.rnn(x, h0) # output shape (batch_size, seq_length, hidden_size)RNN)
        output = output[:, -1, :]
        output = self.fc(output)

        return output
</code></pre>
<p>For the analogous <code>RNNCell</code> -</p>
<pre><code class="lang-auto">## make the model - a recurrent network with one hidden layer, and a fully connected output layer
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(RNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.RNNCell(input_size, hidden_size, num_layers, nonlinearity = 'relu')
        self.fc = nn.Linear(hidden_size, num_classes)


    def forward(self, x):
        h = None
        for t in range(x.shape[1]):
            h = self.rnn(x[:,t,:], h)
        output = self.fc(h)
        return output
</code></pre>
<p>I’ve determined that during validation, the memory usage is quite similar, yet during training the <code>RNN</code> utilizes 3.43GB while the <code>RNNCell</code> uses 8.97GB. I’ve also taken a look at the computational graph using <code>torchviz</code> to determine that the <code>RNN</code> generates one that looks like the following:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d3899c5c1bfc3357e67b8f10a3a72d9db1c51d31.png" data-download-href="https://discuss.pytorch.org/uploads/default/d3899c5c1bfc3357e67b8f10a3a72d9db1c51d31" title="output_simpler_RNNCell2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3899c5c1bfc3357e67b8f10a3a72d9db1c51d31_2_690x461.png" alt="output_simpler_RNNCell2" data-base62-sha1="ublzI9rLuQzUHebijzUKNUhy2kh" width="690" height="461" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3899c5c1bfc3357e67b8f10a3a72d9db1c51d31_2_690x461.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3899c5c1bfc3357e67b8f10a3a72d9db1c51d31_2_1035x691.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/3/d3899c5c1bfc3357e67b8f10a3a72d9db1c51d31.png 2x" data-dominant-color="F4F3F0"><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">output_simpler_RNNCell2</span><span class="informations">1152×771 57.7 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><br>
while the <code>RNNCell</code> generates a much more complicated one, which scales with the sequence length:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/5/2589034e54f7c17d784ef9693fc7058df667306f.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/2589034e54f7c17d784ef9693fc7058df667306f" title="output2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/2589034e54f7c17d784ef9693fc7058df667306f_2_448x500.jpeg" alt="output2" data-base62-sha1="5m3dljcdkY0NwbYJ6EGIFU3cMQ7" width="448" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/2589034e54f7c17d784ef9693fc7058df667306f_2_448x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/2589034e54f7c17d784ef9693fc7058df667306f_2_672x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/5/2589034e54f7c17d784ef9693fc7058df667306f_2_896x1000.jpeg 2x" data-dominant-color="FCFCFB"><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">output2</span><span class="informations">1920×2141 143 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>I know that the <code>RNN()</code> class itself is built on <code>RNNCell()</code>, so I assume some kind of optimization is being done behind the scenes to avoid massive memory usage. I’m aware that PyTorch allows for <a href="https://pytorch.org/docs/stable/checkpoint.html" rel="noopener nofollow ugc">gradient checkpointing</a> and am wondering if this might be what’s happening within the actual <code>RNN</code> code.</p>
<p>So I guess my questions are:</p>
<ol>
<li>How does <code>RNN</code> manage to be memory efficient? I don’t believe it scales with sequence length, but I haven’t tested this extensively myself.</li>
<li>How can I make my code more efficient/as efficient as what is possible with <code>RNN</code>? At the current moment, I am limited on how large my networks can be, even when running them on A100s.</li>
</ol>
<p>Thanks in advance for any and all feedback!</p>

ptrblck: <aside class="quote no-group" data-username="Matthew_Keaton" data-post="1" data-topic="206404">
<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/matthew_keaton/48/55354_2.png" class="avatar"> Matthew_Keaton:</div>
<blockquote>
<p>I know that the <code>RNN()</code> class itself is built on <code>RNNCell()</code>, so I assume some kind of optimization is being done behind the scenes to avoid massive memory usage.</p>
</blockquote>
</aside>
<p>The <code>nn.RNN</code> implementation should use cuDNN while I would guess the <code>RNNCell</code> implementation is using a native implementation and might thus create a lot of intermediates (the second picture is unreadable, so unsure what’s called).</p>

Matthew_Keaton: <p>Ah yes, it looks like the image got compressed. I’ve attached some screenshots, but as you’ve already said it does look like RNN utilizes cuDNN while RNNCell does not. I will look into enabling this for my approach. Thank you!</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/e/3e3e9f8fbda8ada58a08ec19474589f7beafb419.png" data-download-href="https://discuss.pytorch.org/uploads/default/3e3e9f8fbda8ada58a08ec19474589f7beafb419" title="rnncell_zoomed_3"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3e3e9f8fbda8ada58a08ec19474589f7beafb419_2_690x238.png" alt="rnncell_zoomed_3" data-base62-sha1="8SDLVbCA6gXlaTqqvTbP9LQjTrH" width="690" height="238" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3e3e9f8fbda8ada58a08ec19474589f7beafb419_2_690x238.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3e3e9f8fbda8ada58a08ec19474589f7beafb419_2_1035x357.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/e/3e3e9f8fbda8ada58a08ec19474589f7beafb419_2_1380x476.png 2x" data-dominant-color="E8E8E6"><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">rnncell_zoomed_3</span><span class="informations">2875×992 323 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><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/5/a5539b256c3f31d91aaa385879a3d67a7f0987b9.png" data-download-href="https://discuss.pytorch.org/uploads/default/a5539b256c3f31d91aaa385879a3d67a7f0987b9" title="rnncell_zoomed_2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5539b256c3f31d91aaa385879a3d67a7f0987b9_2_647x500.png" alt="rnncell_zoomed_2" data-base62-sha1="nAxUhlqE5WQGdtPW3ZbD5fdvrzr" width="647" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5539b256c3f31d91aaa385879a3d67a7f0987b9_2_647x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5539b256c3f31d91aaa385879a3d67a7f0987b9_2_970x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5539b256c3f31d91aaa385879a3d67a7f0987b9_2_1294x1000.png 2x" data-dominant-color="F6F5F1"><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">rnncell_zoomed_2</span><span class="informations">2049×1583 170 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><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/2/e268334780013c7a8b3e48d5f0be4eb6d61068ea.png" data-download-href="https://discuss.pytorch.org/uploads/default/e268334780013c7a8b3e48d5f0be4eb6d61068ea" title="rnncell_zoomed_1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e268334780013c7a8b3e48d5f0be4eb6d61068ea_2_364x499.png" alt="rnncell_zoomed_1" data-base62-sha1="wiT9HpJocRYbfmwAfuGswRPts7w" width="364" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e268334780013c7a8b3e48d5f0be4eb6d61068ea_2_364x499.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e268334780013c7a8b3e48d5f0be4eb6d61068ea_2_546x748.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e268334780013c7a8b3e48d5f0be4eb6d61068ea_2_728x998.png 2x" data-dominant-color="EFEEE9"><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">rnncell_zoomed_1</span><span class="informations">1272×1746 123 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>