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
For new visitors to this question, please note that torchsummary is now renamed to torchinfo in PIP: https://pypi.org/project/torchinfo/. — QAH, Aug 12 '23 at 03:38

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

21

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,618
10
11

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

Toon 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 5 · Answer 4 · answered Feb 16 '22 at 16:17

5

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

answered Feb 16 '22 at 16:17

minggli

129
2
5

score 3 · Answer 5 · 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

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 0 · Answer 7 · answered May 11 '23 at 07:57

This is a function I wrote:

total_output_elements = 0
def calc_total_activation_size(model, call_the_network_function):
    global total_output_elements
    total_output_elements = 0

    def hook(module, input, output):
        global total_output_elements
        total_output_elements += output.numel()
        
    handle = torch.nn.modules.module.register_module_forward_hook(hook)
    result = call_the_network_function()
    handle.remove()
    return result, total_output_elements

In the hook, you can print the output.shape

The parameters: model - your model call_the_network_function - a function that calls the forward pass of your model, and return the result

score 0 · Answer 8 · answered Jul 27 '23 at 22:27

Inspire by @minggli answer, I used the model's children.

This method will have some steps to modify if not all of the steps are actually in the model's children (e.g. in the ex below a torch.flatten call is in the ResNet18 model's forward method but not in the model's children list).

device = "cuda" # if you want to put on gpu
model = torchvision_models.resnet18(weights="IMAGENET1K_V1")
model.to(device)
batch_size = 4
n_bands = 3
n_rows = 224
n_cols = 224
ex_input = torch.ones((batch_size, n_bands, n_rows, n_cols), device=device)

for i,c in enumerate(list(model.children())):
    # per source code it looks like there is a flatten call which is not in 
    # model.children() & so will need to be added here to get this method to work
    if i == 0:
        layer_output = c(ex_input)
    else:
        if c.__str__() == 'Linear(in_features=512, out_features=10, bias=True)':
            print('layer found')
            print("\t before flatten", layer_output.shape)
            layer_output = torch.flatten(layer_output, 1)
            print("\t after flatten", layer_output.shape)

        layer_output = c(layer_output)

    print(f"children index {i}: {c}")
    print("\t",layer_output.shape)
    print()

score -1 · Answer 9 · 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

11

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

9 Answers9