Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Compute pytorch network layer output size given an input.
#!/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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.