Do pooling layers use memory?

username: agaldran
cooked: <p>Hi,</p>
<p>I would like to ask if using pooling operations in the model definition results in extra memory usage. Imagine I define a simple set of operations as follows:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, pool=True):
        super(ConvBlock, self).__init__()
        self.pool = pool
        blocks = []
        
        blocks.append(nn.Conv2d(in_channels, out_channels, kernel_size=3))
        blocks.append(nn.ReLU())
        blocks.append(nn.BatchNorm2d(out_channels))
        
        if self.pool:
            blocks.append(nn.MaxPool2d(kernel_size = 2))
            
        self.blocks = nn.Sequential(*blocks)
    
    def forward(self, x):
        out = self.blocks(x)
        if not self.pool:
            return F.max_pool2d(out, kernel_size=2)
        return out

conv_block_with_pool = ConvBlock(64,128, pool=True)
conv_block_no_pool = ConvBlock(64,128, pool=False)
</code></pre>
<p>It seems clear that both blocks are doing the same series of operations, only <code>conv_block_no_pool</code> applies pooling in the <code>forward</code> method rather than having it built on the definition.</p>
<p>I’m using <a href="https://github.com/sksq96/pytorch-summary" rel="noopener nofollow ugc">torchsummary</a> to measure memory usage. When using <code>conv_block_with_pool</code>, I obtain the following output:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/7/a7ee55fa2d63694893889ee92cedabf3fdc665e1.png" alt="borrar" data-base62-sha1="nXAmVnwvTVAtYDzkcpr9DXYdkOZ" width="634" height="333"></p>
<p>And when using <code>conv_block_no_pool</code> I get:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/f/b/fb8d6b3acfabd945a84cf5fbb65a28749aaa5bce.png" alt="borrar2" data-base62-sha1="zTkPGnF0wOk8SZizGIUV3a3xX30" width="638" height="317"></p>
<p>Now, I am not sure if <code>torchsummary</code> is not measuring the whole thing in the second case, but I have read <a href="https://www.sicara.ai/blog/2019-28-10-deep-learning-memory-usage-and-pytorch-optimization-tricks" rel="noopener nofollow ugc">elsewhere</a> that even if a layer has no parameters to be learned it still can take over some memory. Is this the case? And if so, should I always prefer/try to add this layers to my models in the forward method, rather than in the model definition itself?</p>
<p>Thanks,</p>
<p>Adrian</p>

username: albanD
cooked: <p>Hi,</p>
<p>Yes pooling will take memory. But this is not the layer itself or its definition that uses the memory, but the fact that it introduces more intermediary results (that are then required to compute the gradients).<br>
If you replaced the pooling with <code> out = out + 1</code>, you will get similar thing.</p>
<p>The pooling might help though as the output will be smaller and the intermediary results of the next layers will be smaller as well.</p>

username: agaldran
cooked: <p>Hi!</p>
<p>Thank you very much for your quick reply. Yes, I understand that, my point would rather be this question:</p>
<blockquote>
<p>Does moving the pooling layer from the model definition (in the constructor) to the forward method bring any advantage in terms of memory consumption?</p>
</blockquote>
<p>I guess this is probably a silly question (and the answer is no). I was just asking because if you look at the output from <code>torchsummary</code> that I was showing in my previous post, this would mean that the memory estimates that <code>torchsummary</code> gives are misleading in this case, and the memory consumed by both “models” would be exactly the same?</p>
<p>Thanks!</p>
<p>Adrian</p>
<p>P.S.: While I have you here, may I also ask if there is any simple way of measuring/estimating the memory spent by a model in Pytorch while doing a forward/backward pass? Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

username: albanD
cooked: <blockquote>
<p>Does moving the pooling layer from the model definition (in the constructor) to the forward method bring any advantage in terms of memory consumption?</p>
</blockquote>
<p>The answer is indeed no.<br>
I’m not familiar how torchsummary works, bu it might indeed be an artefact of it.</p>
<blockquote>
<p>While I have you here, may I also ask if there is any simple way of measuring/estimating the memory spent by a model in Pytorch while doing a forward/backward pass?</p>
</blockquote>
<p>Whatever is your favorite package to measure cpu memory usage in python.<br>
For CPU code, I would recommend python <a href="https://pypi.org/project/memory-profiler/">memory-profiler</a> as it gives very nice outputs.<br>
For GPU code, it is trickier because of both the asynchronous nature of cuda ops and the caching allocator we use. In that case, you can use our own functions to get the used GPU memory and cuda synchronization to get readings manually at different points of your code.</p>