# Function that calculates the output shape of a block

Hi,

is there a function on pytorch that allows to get the shape of an output of a Sequential module knowing the input shape ? For example for this mini block :

``````self.localization = nn.Sequential(
nn.Conv2d(in_channels, 8, kernel_size=7),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True),
nn.Conv2d(8, out_channels, kernel_size=5),
nn.MaxPool2d(2, stride=2),
nn.ReLU(True)
)
``````

We know of course the output formula of a 2D convolution or a max pooling that we can code easily :

``````
def conv_output_shape(h_w, kernel_size=1, stride=1, pad=0, dilation=1):
from math import floor
if type(kernel_size) is not tuple:
kernel_size = (kernel_size, kernel_size)
h = floor(((h_w[0] + (2 * pad) - (dilation * (kernel_size[0] - 1)) - 1) / stride) + 1)
w = floor(((h_w[1] + (2 * pad) - (dilation * (kernel_size[1] - 1)) - 1) / stride) + 1)
return h, w
``````

But as I said Iām looking to do the same thing so a function I can call in a network for a group of operations so for a nn.Sequential or even a network in general.

You can attach forward hook to get the output size of certain module.
Example:

``````import torch
import torch.nn as nn

seq = nn.Sequential(nn.Conv2d(1, 2, 3), nn.BatchNorm2d(2), nn.Conv2d(2, 3, 3))

def forward_hook(module, input, output):
print(output.size())

seq.register_forward_hook(forward_hook)
seq(torch.rand(1, 1, 10, 10))
``````

which prints out torch.Size([1, 3, 6, 6]).

1 Like

@sio277 wah great !! thanks