Compute pytorch network layer output size given an input.

pytorch_compute_out_size.py

#!/usr/bin/env python

"""
Compute pytorch network layer output size given an input.
"""

import numpy as np
import torch
import torch.autograd as autograd
import torch.nn as nn

def compute_out_size(in_size, mod):
    """
    Compute output size of Module `mod` given an input with size `in_size`.
    """
    
    f = mod.forward(autograd.Variable(torch.Tensor(1, *in_size)))
    return int(np.prod(f.size()[1:]))

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.m = nn.Sequential(nn.Conv2d(1, 5, 3))
                    
    def forward(self, x):
        return self.m.forward(x)

net = Network()
x = torch.Tensor(1, 1, 100, 100) # shape = (batch size, channels, height, width)
print compute_size(x.size(), net.forward)