Hi,
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:
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)

It seems clear that both blocks are doing the same series of operations, only conv_block_no_pool applies pooling in the forward method rather than having it built on the definition.
I’m using torchsummary 12 to measure memory usage. When using conv_block_with_pool, I obtain the following output:

And when using conv_block_no_pool I get:

Now, I am not sure if torchsummary is not measuring the whole thing in the second case, but I have read elsewhere 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?
Thanks,
Adrian