Function that calculates the output shape of a block

vision
1

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.

2

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
3

@sio277 wah great !! thanks :wink: