torchinfo 来打印网络

最新推荐文章于 2024-06-04 10:45:52 发布

Zou验钞

最新推荐文章于 2024-06-04 10:45:52 发布

阅读量271

点赞数

文章标签：人工智能

本文链接：https://blog.csdn.net/zycnice/article/details/128193044

版权

使用torchinfo 来打印网络的每一层尺寸等信息

from cascadenet import U_Network
from torchinfo import summary
import torchvision

model = U_Network()
summary(model, (1,2, 64, 64, 64), depth=3)

结果是

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
U_Network                                --                        --
├─ModuleList: 1-1                        --                        --
├─ModuleList: 1-2                        --                        --
├─ModuleList: 1-1                        --                        --
│    └─Sequential: 2-1                   [1, 16, 64, 64, 64]       --
│    │    └─Conv3d: 3-1                  [1, 16, 64, 64, 64]       880
│    │    └─InstanceNorm3d: 3-2          [1, 16, 64, 64, 64]       --
│    │    └─LeakyReLU: 3-3               [1, 16, 64, 64, 64]       --
├─MaxPool3d: 1-3                         [1, 16, 32, 32, 32]       --
├─ModuleList: 1-1                        --                        --
│    └─Sequential: 2-2                   [1, 32, 32, 32, 32]       --
│    │    └─Conv3d: 3-4                  [1, 32, 32, 32, 32]       13,856
│    │    └─InstanceNorm3d: 3-5          [1, 32, 32, 32, 32]       --
│    │    └─LeakyReLU: 3-6               [1, 32, 32, 32, 32]       --
├─MaxPool3d: 1-4                         [1, 32, 16, 16, 16]       --
├─ModuleList: 1-1                        --                        --
│    └─Sequential: 2-3                   [1, 32, 16, 16, 16]       --
│    │    └─Conv3d: 3-7                  [1, 32, 16, 16, 16]       27,680
│    │    └─InstanceNorm3d: 3-8          [1, 32, 16, 16, 16]       --
│    │    └─LeakyReLU: 3-9               [1, 32, 16, 16, 16]       --
├─MaxPool3d: 1-5                         [1, 32, 8, 8, 8]          --
├─ModuleList: 1-1                        --                        --
│    └─Sequential: 2-4                   [1, 32, 8, 8, 8]          --
│    │    └─Conv3d: 3-10                 [1, 32, 8, 8, 8]          27,680
│    │    └─InstanceNorm3d: 3-11         [1, 32, 8, 8, 8]          --
│    │    └─LeakyReLU: 3-12              [1, 32, 8, 8, 8]          --
├─MaxPool3d: 1-6                         [1, 32, 4, 4, 4]          --
├─ModuleList: 1-1                        --                        --
│    └─Sequential: 2-5                   [1, 32, 4, 4, 4]          --
│    │    └─Conv3d: 3-13                 [1, 32, 4, 4, 4]          27,680
│    │    └─InstanceNorm3d: 3-14         [1, 32, 4, 4, 4]          --
│    │    └─LeakyReLU: 3-15              [1, 32, 4, 4, 4]          --
├─Upsample: 1-7                          [1, 32, 8, 8, 8]          --
├─ModuleList: 1-2                        --                        --
│    └─Sequential: 2-6                   [1, 32, 8, 8, 8]          --
│    │    └─Conv3d: 3-16                 [1, 32, 8, 8, 8]          27,680
│    │    └─InstanceNorm3d: 3-17         [1, 32, 8, 8, 8]          --
│    │    └─LeakyReLU: 3-18              [1, 32, 8, 8, 8]          --
│    └─Sequential: 2-7                   [1, 32, 8, 8, 8]          --
│    │    └─Conv3d: 3-19                 [1, 32, 8, 8, 8]          55,328
│    │    └─InstanceNorm3d: 3-20         [1, 32, 8, 8, 8]          --
│    │    └─LeakyReLU: 3-21              [1, 32, 8, 8, 8]          --
├─Upsample: 1-8                          [1, 32, 16, 16, 16]       --
├─ModuleList: 1-2                        --                        --
│    └─Sequential: 2-8                   [1, 32, 16, 16, 16]       --
│    │    └─Conv3d: 3-22                 [1, 32, 16, 16, 16]       55,328
│    │    └─InstanceNorm3d: 3-23         [1, 32, 16, 16, 16]       --
│    │    └─LeakyReLU: 3-24              [1, 32, 16, 16, 16]       --
├─Upsample: 1-9                          [1, 32, 32, 32, 32]       --
├─ModuleList: 1-2                        --                        --
│    └─Sequential: 2-9                   [1, 32, 32, 32, 32]       --
│    │    └─Conv3d: 3-25                 [1, 32, 32, 32, 32]       55,328
│    │    └─InstanceNorm3d: 3-26         [1, 32, 32, 32, 32]       --
│    │    └─LeakyReLU: 3-27              [1, 32, 32, 32, 32]       --
├─Upsample: 1-10                         [1, 32, 64, 64, 64]       --
├─ModuleList: 1-2                        --                        --
│    └─Sequential: 2-10                  [1, 32, 64, 64, 64]       --
│    │    └─Conv3d: 3-28                 [1, 32, 64, 64, 64]       41,504
│    │    └─InstanceNorm3d: 3-29         [1, 32, 64, 64, 64]       --
│    │    └─LeakyReLU: 3-30              [1, 32, 64, 64, 64]       --
├─Sequential: 1-11                       [1, 16, 64, 64, 64]       --
│    └─Conv3d: 2-11                      [1, 16, 64, 64, 64]       13,840
│    └─InstanceNorm3d: 2-12              [1, 16, 64, 64, 64]       --
│    └─LeakyReLU: 2-13                   [1, 16, 64, 64, 64]       --
├─Sequential: 1-12                       [1, 8, 64, 64, 64]        --
│    └─Conv3d: 2-14                      [1, 8, 64, 64, 64]        3,464
│    └─InstanceNorm3d: 2-15              [1, 8, 64, 64, 64]        --
│    └─LeakyReLU: 2-16                   [1, 8, 64, 64, 64]        --
├─Conv3d: 1-13                           [1, 3, 64, 64, 64]        651
==========================================================================================
Total params: 350,899
Trainable params: 350,899
Non-trainable params: 0
Total mult-adds (G): 18.48
==========================================================================================
Input size (MB): 2.10
Forward/backward pass size (MB): 176.57
Params size (MB): 1.40
Estimated Total Size (MB): 180.07
==========================================================================================

Process finished with exit code 0