How to get an output dimension for each layer of the Neural Network in Pytorch?

Question

class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    self.net = nn.Sequential(
      nn.Conv2d(in_channels = 3, out_channels = 16), 
      nn.ReLU(), 
      nn.MaxPool2d(2),
      nn.Conv2d(in_channels = 16, out_channels = 16), 
      nn.ReLU(),
      Flatten(),
      nn.Linear(4096, 64),
      nn.ReLU(),
      nn.Linear(64, 10))

  def forward(self, x):
    return self.net(x)

I have created this model without a firm knowledge in Neural Network and I just fixed parameters until it worked in the training. I am not sure how to get the output dimension for each layer (e.g. output dimension after the first layer).

Is there an easy way to do this in Pytorch?

Does this answer your question? [Model summary in pytorch](https://stackoverflow.com/questions/42480111/model-summary-in-pytorch) — iacob, Mar 15 '21 at 15:50

Idan Azuri · Answer 1 · 2019-04-27T12:58:06.417

You can use torchsummary, for instance, for ImageNet dimension(3x224x224):

from torchvision import models
from torchsummary import summary

vgg = models.vgg16()
summary(vgg, (3, 224, 224)


----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 224, 224]           1,792
              ReLU-2         [-1, 64, 224, 224]               0
            Conv2d-3         [-1, 64, 224, 224]          36,928
              ReLU-4         [-1, 64, 224, 224]               0
         MaxPool2d-5         [-1, 64, 112, 112]               0
            Conv2d-6        [-1, 128, 112, 112]          73,856
              ReLU-7        [-1, 128, 112, 112]               0
            Conv2d-8        [-1, 128, 112, 112]         147,584
              ReLU-9        [-1, 128, 112, 112]               0
        MaxPool2d-10          [-1, 128, 56, 56]               0
           Conv2d-11          [-1, 256, 56, 56]         295,168
             ReLU-12          [-1, 256, 56, 56]               0
           Conv2d-13          [-1, 256, 56, 56]         590,080
             ReLU-14          [-1, 256, 56, 56]               0
           Conv2d-15          [-1, 256, 56, 56]         590,080
             ReLU-16          [-1, 256, 56, 56]               0
        MaxPool2d-17          [-1, 256, 28, 28]               0
           Conv2d-18          [-1, 512, 28, 28]       1,180,160
             ReLU-19          [-1, 512, 28, 28]               0
           Conv2d-20          [-1, 512, 28, 28]       2,359,808
             ReLU-21          [-1, 512, 28, 28]               0
           Conv2d-22          [-1, 512, 28, 28]       2,359,808
             ReLU-23          [-1, 512, 28, 28]               0
        MaxPool2d-24          [-1, 512, 14, 14]               0
           Conv2d-25          [-1, 512, 14, 14]       2,359,808
             ReLU-26          [-1, 512, 14, 14]               0
           Conv2d-27          [-1, 512, 14, 14]       2,359,808
             ReLU-28          [-1, 512, 14, 14]               0
           Conv2d-29          [-1, 512, 14, 14]       2,359,808
             ReLU-30          [-1, 512, 14, 14]               0
        MaxPool2d-31            [-1, 512, 7, 7]               0
           Linear-32                 [-1, 4096]     102,764,544
             ReLU-33                 [-1, 4096]               0
          Dropout-34                 [-1, 4096]               0
           Linear-35                 [-1, 4096]      16,781,312
             ReLU-36                 [-1, 4096]               0
          Dropout-37                 [-1, 4096]               0
           Linear-38                 [-1, 1000]       4,097,000
================================================================
Total params: 138,357,544
Trainable params: 138,357,544
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 218.59
Params size (MB): 527.79
Estimated Total Size (MB): 746.96
----------------------------------------------------------------

Source: model-summary-in-pytorch

@Dawn17 It's the dimensions of a single image (for MNIST it is 1x28x28) — Idan Azuri, Apr 27 '19 at 16:55
I am getting this error `RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 3, 3], but got 5-dimensional input of size [2, 1, 3, 224, 224] instead` What size do I have to put? — Dawn17, Apr 29 '19 at 01:12
@Dawn17 I need to see your code to assist you, but I guess that you run in your network MNIST which is 1x28x28 and VGG input is 3x224x224. So, first in the forward method try to reshape it like: 'out.view(out.shape[0], -1)' and second, change the model to yours instead of the VGG in my example. — Idan Azuri, Apr 29 '19 at 12:07

David Ng · Accepted Answer · 2019-04-27T09:35:26.483

14

A simple way is:

Pass the input to the model.
Print the size of the output after passing every layer.

class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    self.net = nn.Sequential(
      nn.Conv2d(in_channels = 3, out_channels = 16), 
      nn.ReLU(), 
      nn.MaxPool2d(2),
      nn.Conv2d(in_channels = 16, out_channels = 16), 
      nn.ReLU(),
      Flatten(),
      nn.Linear(4096, 64),
      nn.ReLU(),
      nn.Linear(64, 10))

  def forward(self, x):
    for layer in self.net:
        x = layer(x)
        print(x.size())
    return x

model = Model()
x = torch.randn(1, 3, 224, 224)

# Let's print it
model(x)

But be careful with the input size because you are using nn.Linear in your net. It would cause incompatible input size for nn.Linear if your input size is not 4096.

edited Apr 27 '19 at 09:35

answered Apr 27 '19 at 05:29

David Ng

1,468
9
10

What's the (1, 3, 244, 244) ? – Dawn17 Apr 27 '19 at 12:20
just create a pseudo example of input: `x = torch.randn(1, 3, 244, 244)` – David Ng Apr 27 '19 at 14:29
what i have multiple networks in my model and multiple inputs and some of the inputs look something like torch.Size([20,16000])? – monart Oct 14 '20 at 14:51

Tung Tran · Answer 3 · 2020-06-04T14:28:22.710

Like David Ng's answer but a tad shorter:

def get_output_shape(model, image_dim):
    return model(torch.rand(*(image_dim))).data.shape

In this example I needed to figure out the input of the last Linear layer:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.expected_input_shape = (1, 1, 192, 168)
        self.conv1 = nn.Conv2d(1, 32, 3, 1) 
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.maxpool1 = nn.MaxPool2d(2)
        self.maxpool2 = nn.MaxPool2d(3)

        # Calculate the input of the Linear layer
        conv1_out = get_output_shape(self.maxpool1, get_output_shape(conv1, self.expected_input_shape))
        conv2_out = get_output_shape(self.maxpool2, get_output_shape(conv2, conv1_out)) 
        fc1_in = np.prod(list(conv2_out)) # Flatten

        self.fc1 = nn.Linear(fc1_in, 38)

    def forward(self, x):
        x = self.conv1(x) 
        x = F.relu(x)
        x = self.maxpool1(x) 
        x = self.conv2(x)
        x = F.relu(x)
        x = self.maxpool2(x) 
        x = self.dropout1(x) 
        x = torch.flatten(x, 1) # flatten to a single dimension
        x = self.fc1(x) 
        output = F.log_softmax(x, dim=1) 
        return output

This way, if I make changes to previous layers, I won't have to calculate all over again!

My answer is based on this answer

score 3 · Answer 4 · answered Nov 03 '20 at 21:35

Another way to get the size after a certain layer in an nn.Sequential container is to add a custom Module that just prints out the size of the input.

class PrintSize(nn.Module):
  def __init__(self):
    super(PrintSize, self).__init__()
    
  def forward(self, x):
    print(x.shape)
    return x

And now you can do:

model = nn.Sequential(
    nn.Conv2d(3, 10, 5, 1),
    // lots of convolutions, pooling, etc.
    nn.Flatten(),
    PrintSize(),
    nn.Linear(1, 12), // the input dim of 1 is just a placeholder
)

Now, you can do model(x) and it will print out the shape of the output after the Conv2d layer ran. This is useful if you have a lot of convolutions and want to figure out what the final dimensions are for the first fully connected layer. You don't need to reformat your nn.Sequential as a Module and can just drop in this helper class with one-line.

If we do not want to change existing model then layer.register_forward_hook is better — x4444, Apr 23 '21 at 04:51

score 2 · Answer 5 · answered Feb 16 '22 at 16:17

2

for layer in model.children():
    if hasattr(layer, 'out_features'):
        print(layer.out_features)

answered Feb 16 '22 at 16:17

minggli

89
1
5

Jose Solorzano · Answer 6 · 2021-03-17T17:34:27.970

Here's a solution in the form of a helper function:

def get_tensor_dimensions_impl(model, layer, image_size, for_input=False):
    t_dims = None
    def _local_hook(_, _input, _output):
        nonlocal t_dims
        t_dims = _input[0].size() if for_input else _output.size()
        return _output    
    layer.register_forward_hook(_local_hook)
    dummy_var = torch.zeros(1, 3, image_size, image_size)
    model(dummy_var)
    return t_dims

Example:

from torchvision import models, transforms

a_model = models.squeezenet1_0(pretrained=True) 
get_tensor_dimensions_impl(a_model, a_model._modules['classifier'], 224)

Output is:

torch.Size([1, 1000, 1, 1])

score -1 · Answer 7 · answered Jun 03 '21 at 08:18

-1

Maybe you can try print(model.state_dict()['next_layer.weight'].shape). This gives you a hint of the output shape from last layer.

answered Jun 03 '21 at 08:18

James.Q.Freeman

1

How to get an output dimension for each layer of the Neural Network in Pytorch?

7 Answers7