Print(model) needs output shapes

Oktai15 · April 2, 2018, 8:41am

Hi there!

I think that everyone wants something in PyTorch like model.summary() in Keras.
Even people wrote it themselves, for examples:

gist.github.com

https://gist.github.com/HTLife/b6640af9d6e7d765411f8aa9aa94b837

summary.py

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

This file has been truncated. show original

But PyTorch team doesn’t do it… What is the reason for it?