Implement similar PyTorch function as model.summary() in keras?

[iabhi7]

model.summary in keras gives a very fine visualization of your model and it's very convenient when it comes to debugging the network. Can we try to implement something like it in PyTorch?

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @albanD @mruberry

[apaszke]

What kind of information that's not in str(model) would you like to see? Output shapes?

[ncullen93]

Sure you can! Here's something to get you started. (Adapted from other code, so it's not tested in the wild). Note that you HAVE to know the input size, and you HAVE to make a forward pass through the network. Those are the only reqs I think.

    def summary(input_size, model):
        def register_hook(module):
            def hook(module, input, output):
                class_name = str(module.__class__).split('.')[-1].split("'")[0]
                module_idx = len(summary)

                m_key = '%s-%i' % (class_name, module_idx+1)
                summary[m_key] = OrderedDict()
                summary[m_key]['input_shape'] = list(input[0].size())
                summary[m_key]['input_shape'][0] = -1
                summary[m_key]['output_shape'] = list(output.size())
                summary[m_key]['output_shape'][0] = -1

                params = 0
                if hasattr(module, 'weight'):
                    params += th.prod(th.LongTensor(list(module.weight.size())))
                    if module.weight.requires_grad:
                        summary[m_key]['trainable'] = True
                    else:
                        summary[m_key]['trainable'] = False
                if hasattr(module, 'bias'):
                    params +=  th.prod(th.LongTensor(list(module.bias.size())))
                summary[m_key]['nb_params'] = params
                
            if not isinstance(module, nn.Sequential) and \
               not isinstance(module, nn.ModuleList) and \
               not (module == model):
                hooks.append(module.register_forward_hook(hook))
        
        # check if there are multiple inputs to the network
        if isinstance(input_size[0], (list, tuple)):
            x = [Variable(th.rand(1,*in_size)) for in_size in input_size]
        else:
            x = Variable(th.rand(1,*input_size))

        # create properties
        summary = OrderedDict()
        hooks = []
        # register hook
        model.apply(register_hook)
        # make a forward pass
        model(x)
        # remove these hooks
        for h in hooks:
            h.remove()

        return summary

Here's an example of Keras summary, FYI:

[iabhi7]

@apaszke str(output) works well for most cases but missed out on some output_shape is one of them
@ncullen93 The code looks/works well. Can we think about a PR with some code refactoring to this?

[aditya1702]

Hello, I am new to pytorch and contributing to it. Can I try this one out?

[iabhi7]

@aditya1702 yeah sure, as my information goes till now we don't have any functions to see the output_size and something similiar to keras model.summary

[EtienneDesticourt]

I think it would be useful to add hooks to the output of str(model).
I use str(model) in my logging system and currently I can't know whether my models used weight normalizations or not since it's a hook.

[isaykatsman]

After @ncullen93 posted his code, I added model summary to my local build of pytorch. Recently decided to clean the code up a little and make a PR. Can hopefully merge soon.

[iabhi7]

Hi @isaykatsman the implementation is nice. Thanks.

[HTLife]

@ncullen93
It that 'Variable' at the line x = Variable(th.rand(1,*input_size)) a package?
How can I deal with this error?

global name 'Variable' is not defined

[apaszke]

from torch.autograd import Variable

[HTLife]

I did some modification on the print out style to fit the style of Keras. Hope it helps.

(2018/05/18 update)
sksq96 re-organized the code into python package sksq96/pytorch-summary

[backpropper]

This doesn't work for LSTMs where the output is a tuple.

[ShuvenduRoy]

@ncullen93 It does not work with sequential() function. Do you have a solution for this?

class Discriminator(nn.Module):
    def __init__(self, in_channels=3):
        super(Discriminator, self).__init__()

        def discriminator_block(in_filters, out_filters, normalize=True):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
            if normalize:
                layers.append(nn.InstanceNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *discriminator_block(in_channels, 64, normalize=False),
            *discriminator_block(64, 128),
            *discriminator_block(128, 256),
            *discriminator_block(256, 512),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, 4, padding=1)
        )

    def forward(self, img):
        return self.model(img)

Traceback (most recent call last):
  File ".\test.py", line 118, in <module>
    summary(model, (3, 28, 28))
  File "C:\Users\bikas\Anaconda3\lib\site-packages\torchsummary\torchsummary.py", line 57, in summary
    model(x)
  File "C:\Users\bikas\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File ".\test.py", line 93, in forward
    return self.model(img)
  File "C:\Users\bikas\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\bikas\Anaconda3\lib\site-packages\torch\nn\modules\container.py", line 91, in forward
    input = module(input)
  File "C:\Users\bikas\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 493, in __call__
    hook_result = hook(self, input, result)
  File "C:\Users\bikas\Anaconda3\lib\site-packages\torchsummary\torchsummary.py", line 26, in hook
    params += torch.prod(torch.LongTensor(list(module.weight.size())))
AttributeError: 'NoneType' object has no attribute 'size'

[jongwook]

I like the tree format of the builtin __repr__() but wanted to show the number of parameters, like:

import sys
from functools import reduce

from torch.nn.modules.module import _addindent

def summary(model, file=sys.stderr):
    def repr(model):
        # We treat the extra repr like the sub-module, one item per line
        extra_lines = []
        extra_repr = model.extra_repr()
        # empty string will be split into list ['']
        if extra_repr:
            extra_lines = extra_repr.split('\n')
        child_lines = []
        total_params = 0
        for key, module in model._modules.items():
            mod_str, num_params = repr(module)
            mod_str = _addindent(mod_str, 2)
            child_lines.append('(' + key + '): ' + mod_str)
            total_params += num_params
        lines = extra_lines + child_lines

        for name, p in model._parameters.items():
            total_params += reduce(lambda x, y: x * y, p.shape)

        main_str = model._get_name() + '('
        if lines:
            # simple one-liner info, which most builtin Modules will use
            if len(extra_lines) == 1 and not child_lines:
                main_str += extra_lines[0]
            else:
                main_str += '\n  ' + '\n  '.join(lines) + '\n'

        main_str += ')'
        if file is sys.stderr:
            main_str += ', \033[92m{:,}\033[0m params'.format(total_params)
        else:
            main_str += ', {:,} params'.format(total_params)
        return main_str, total_params

    string, count = repr(model)
    if file is not None:
        print(string, file=file)
    return count

[PistonY]

What kind of information that's not in str(model) would you like to see? Output shapes?

Total parameters and total FLOPS.

[ray-lee-94]

Can i draw the forward pass graph similar to torchviz ?
If i summary the model and get the visulization, that would be good.

[zamnius]

NameError: name 'th' is not defined
For example in the line: x = [Variable(th.rand(1,*in_size)).type(dtype) for in_size in input_size]
What is missing?

[hmrishavbandy]

@zamnius , here are the complete list of imports:

from torch.autograd import Variable
import torch as th
from torch import nn as nn
from collections import OrderedDict
from model import UNet3D

@ShuvenduBikash , this is probably what you figured out:
The lines 15-22 should be like this as all layers dont have bias.

                if hasattr(module,'weight') and module.weight is not None:
                    params += th.prod(th.LongTensor(list(module.weight.size())))
                    if module.weight.requires_grad:
                        summary[m_key]['trainable'] = True
                    else:
                        summary[m_key]['trainable'] = False
                if hasattr(module,'bias') and module.bias is not None:
                    params +=  th.prod(th.LongTensor(list(module.bias.size())))

[zou3519]

I don't remember why this was originally in the low-priority bin, but I'm tentatively marking this as triage review for discussion on whether or not this should be high-pri

[jbschlosser]

Removing high-pri for now since this is provided by pytorch-summary, which appears to work for most cases described here (including nn.Sequential models).

As of now, I don't see huge benefits to maintaining this within PyTorch core and would prefer to keep it in the separate repo, but I'm open to being convinced otherwise if it doesn't support your use case, etc.

[tuiiitendinh]

Hi, i'm new to pytorch.
Is there anyway to plot the graph of the model as well as export the sumary of the model without initialize an instance of the model? Like Keras, I would just load the model and then visualize it.

[mert-kurttutan]

Hi, i'm new to pytorch. Is there anyway to plot the graph of the model as well as export the sumary of the model without initialize an instance of the model? Like Keras, I would just load the model and then visualize it.

Hi @tuiiitendinh,
I created a package to visualize pytorch models. You can see it, here
There are a few notebooks to get familiar with in addition to documentation.
Any feedback/issue/pr is much appreciated