Is there similar pytorch function as model.summary() as keras?

is there similar pytorch function as model.summary() as keras?

1 Like

repr(model) gives something fairly close.

2 Likes

or easier to remember:

print(model)

gives the same result as repr(model)

2 Likes

Also, if you just want the number of parameters as opposed to the whole model, you can do something like: sum([param.nelement() for param in model.parameters()])

Yes, you can get exact Keras representation, using this code.

Example for VGG16

from torchvision import models
from summary import summary

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

----------------------------------------------------------------
        Layer (type)               Output Shpae         Param #
================================================================
            Conv2d-1         [-1, 64, 224, 224]            1792
              ReLU-2         [-1, 64, 224, 224]               0
            Conv2d-3         [-1, 64, 224, 224]           36928
              ReLU-4         [-1, 64, 224, 224]               0
         MaxPool2d-5         [-1, 64, 112, 112]               0
            Conv2d-6        [-1, 128, 112, 112]           73856
              ReLU-7        [-1, 128, 112, 112]               0
            Conv2d-8        [-1, 128, 112, 112]          147584
              ReLU-9        [-1, 128, 112, 112]               0
        MaxPool2d-10          [-1, 128, 56, 56]               0
           Conv2d-11          [-1, 256, 56, 56]          295168
             ReLU-12          [-1, 256, 56, 56]               0
           Conv2d-13          [-1, 256, 56, 56]          590080
             ReLU-14          [-1, 256, 56, 56]               0
           Conv2d-15          [-1, 256, 56, 56]          590080
             ReLU-16          [-1, 256, 56, 56]               0
        MaxPool2d-17          [-1, 256, 28, 28]               0
           Conv2d-18          [-1, 512, 28, 28]         1180160
             ReLU-19          [-1, 512, 28, 28]               0
           Conv2d-20          [-1, 512, 28, 28]         2359808
             ReLU-21          [-1, 512, 28, 28]               0
           Conv2d-22          [-1, 512, 28, 28]         2359808
             ReLU-23          [-1, 512, 28, 28]               0
        MaxPool2d-24          [-1, 512, 14, 14]               0
           Conv2d-25          [-1, 512, 14, 14]         2359808
             ReLU-26          [-1, 512, 14, 14]               0
           Conv2d-27          [-1, 512, 14, 14]         2359808
             ReLU-28          [-1, 512, 14, 14]               0
           Conv2d-29          [-1, 512, 14, 14]         2359808
             ReLU-30          [-1, 512, 14, 14]               0
        MaxPool2d-31            [-1, 512, 7, 7]               0
           Linear-32                 [-1, 4096]       102764544
             ReLU-33                 [-1, 4096]               0
          Dropout-34                 [-1, 4096]               0
           Linear-35                 [-1, 4096]        16781312
             ReLU-36                 [-1, 4096]               0
          Dropout-37                 [-1, 4096]               0
           Linear-38                 [-1, 1000]         4097000
================================================================
Total params: 138357544
Trainable params: 138357544
Non-trainable params: 0
----------------------------------------------------------------
3 Likes

Actually, there’s a difference between keras model.summary() and print(model) in pytorch. print(model in pytorch only print the layers defined in the init function of the class but not the model architecture defined in forward function. Keras model.summary() actually prints the model architecture with input and output shape along with trainable and non trainable parameters.
I haven’t found anything like that in PyTorch. I end up writing bunch of print statements in forward function to determine the input and output shape.

3 Likes

Is the suggestion from @sksq96 not sufficient for the shape information?

If you check the repo, that summary is only printing what’s in init function not the actual forward function where you will apply batch normalization, relu, maxpooling, global max pooling like stuff.
VGG model is printing because its implemented that way, meaning ReLu layer is defined in init function instead of Functional relu. I faced huge problem when implementing UNet, and that’s just one example.
Check this from the repo

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # PyTorch v0.4.0
model = Net().to(device)

summary(model, (1, 28, 28))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------

This summary absolutely doesn’t make any sense.

The summary would still give you all parameters, no?
If you are interested in printing the graph, have a look at e.g. this topic.

I just want a easy function call to print the model summary the way Keras do. PyTorch should have added that. I should not be doing all kind of tricks just to see my model summary with input and output shapes of every layer. I love working in PyTorch that’s why I am looking for that type of function that would make model development easy.

pytorch-summary should yield the same output as Keras does.
What are you missing?

1 Like

pytorch-summary can’t handle things like lists of tensors in forward(). Pytorch should definitely have a summary() function which can handle any operations done inside the model.

2 Likes

you can try this library pytorch-model-summary. I rewrote to support scenario you described, adding other options as well

This worked for me

!pip install torchsummary
from torchsummary import summary
summary(your_model, input_size=(channels, H, W))
3 Likes

Really love your work ! works perfectly. thank you.

For a given input shape, you can use the torchinfo (formerly torchsummary) package:

Torchinfo provides information complementary to what is provided by print(your_model) in PyTorch, similar to Tensorflow’s model.summary()…

Example:

from torchinfo import summary

model = ConvNet()
batch_size = 16
summary(model, input_size=(batch_size, 1, 28, 28)
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
├─Conv2d (conv1): 1-1                    [5, 10, 24, 24]           260
├─Conv2d (conv2): 1-2                    [5, 20, 8, 8]             5,020
├─Dropout2d (conv2_drop): 1-3            [5, 20, 8, 8]             --
├─Linear (fc1): 1-4                      [5, 50]                   16,050
├─Linear (fc2): 1-5                      [5, 10]                   510
==========================================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
Total mult-adds (M): 7.69
==========================================================================================
Input size (MB): 0.05
Forward/backward pass size (MB): 0.91
Params size (MB): 0.09
Estimated Total Size (MB): 1.05
==========================================================================================

Note: Unlike Keras, PyTorch has a dynamic computational graph which can adapt to any compatible input shape across multiple calls e.g. any sufficiently large image size (for a fully convolutional network).

As such, it cannot present an inherent set of input/output shapes for each layer, as these are input-dependent, and why in the above package you must specify the input dimensions.