Keras style model.summary() in PyTorch - 打印 PyTorch 模型结构

蔚落 2023-02-19 11:25 3阅读 0赞

#  Keras style `model.summary()` in PyTorch - 打印 PyTorch 模型结构 #

[https://github.com/sksq96/pytorch-summary][https_github.com_sksq96_pytorch-summary]  
[https://pypi.org/project/torchsummary/][https_pypi.org_project_torchsummary]

Model summary in PyTorch similar to `model.summary()` in Keras.

Keras has a neat API to view the visualization of the model which is very helpful while debugging your network. Here is a barebone code to try and mimic the same in PyTorch. The aim is to provide information complementary to, what is not provided by `print(your_model)` in PyTorch.  
Keras 有一个简洁的 API 可以查看模型的可视化，这在调试网络时非常有帮助。这是一个尝试在 PyTorch 中模仿的准系统代码。目的是提供在 PyTorch 中的 `print(your_model)` 无法提供的互补的信息。

mimic /ˈmɪmɪk/：vt. 模仿，摹拟 n. 效颦者，模仿者，仿制品，小丑 adj. 模仿的，模拟的，假装的
    barebone /'bɛə'bəun/：n. 准系统，准系统电脑

## 1. usage ##

`pip install torchsummary`

(pt-1.4_py-3.6) yongqiang@yongqiang:~$ pip install torchsummary
    Collecting torchsummary
      Using cached torchsummary-1.5.1-py3-none-any.whl (2.8 kB)
    Installing collected packages: torchsummary
    Successfully installed torchsummary-1.5.1
    (pt-1.4_py-3.6) yongqiang@yongqiang:~$

from torchsummary import summary
    summary(your_model, input_size=(channels, H, W))

Note that the `input_size` is required to make a forward pass through the network.

## 2. CNN for MNIST ##

#!/usr/bin/env python
    # -*- coding: utf-8 -*-
    # yongqiang cheng
    
    from __future__ import absolute_import
    from __future__ import division
    from __future__ import print_function
    
    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
    model = Net().to(device)
    
    summary(model, (1, 28, 28))

## 3. VGG16 ##

#!/usr/bin/env python
    # -*- coding: utf-8 -*-
    # yongqiang cheng
    
    from __future__ import absolute_import
    from __future__ import division
    from __future__ import print_function
    
    import torch
    from torchsummary import summary
    from torchvision import models
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    vgg = models.vgg16().to(device)
    
    summary(vgg, (3, 224, 224))

/home/yongqiang/miniconda3/envs/pt-1.4_py-3.6/bin/python /home/yongqiang/pytorch_work/end2end-asr-pytorch-example/yongqiang.py
    ----------------------------------------------------------------
            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
    AdaptiveAvgPool2d-32            [-1, 512, 7, 7]               0
               Linear-33                 [-1, 4096]     102,764,544
                 ReLU-34                 [-1, 4096]               0
              Dropout-35                 [-1, 4096]               0
               Linear-36                 [-1, 4096]      16,781,312
                 ReLU-37                 [-1, 4096]               0
              Dropout-38                 [-1, 4096]               0
               Linear-39                 [-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.78
    Params size (MB): 527.79
    Estimated Total Size (MB): 747.15
    ----------------------------------------------------------------
    
    Process finished with exit code 0

## 4. Multiple Inputs ##

#!/usr/bin/env python
    # -*- coding: utf-8 -*-
    # yongqiang cheng
    
    from __future__ import absolute_import
    from __future__ import division
    from __future__ import print_function
    
    import torch
    import torch.nn as nn
    from torchsummary import summary
    
    
    class SimpleConv(nn.Module):
        def __init__(self):
            super(SimpleConv, self).__init__()
            self.features = nn.Sequential(
                nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1),
                nn.ReLU(),
            )
    
        def forward(self, x, y):
            x1 = self.features(x)
            x2 = self.features(y)
            return x1, x2
    
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SimpleConv().to(device)
    
    summary(model, [(1, 16, 16), (1, 28, 28)])

## References ##

[http://jkimmel.net/pytorch\_estimating\_model\_size/][http_jkimmel.net_pytorch_estimating_model_size]  
[https://github.com/pytorch/pytorch/issues/2001][https_github.com_pytorch_pytorch_issues_2001]  
[https://github.com/TangShengqin/pytorch\_learn/blob/master/model\_flops.py][https_github.com_TangShengqin_pytorch_learn_blob_master_model_flops.py]

[https_github.com_sksq96_pytorch-summary]: https://github.com/sksq96/pytorch-summary
[https_pypi.org_project_torchsummary]: https://pypi.org/project/torchsummary/
[http_jkimmel.net_pytorch_estimating_model_size]: http://jkimmel.net/pytorch_estimating_model_size/
[https_github.com_pytorch_pytorch_issues_2001]: https://github.com/pytorch/pytorch/issues/2001
[https_github.com_TangShengqin_pytorch_learn_blob_master_model_flops.py]: https://github.com/TangShengqin/pytorch_learn/blob/master/model_flops.py