模型可视化工具 torchinfo 计算每层输出的Output Shape

官方文档,pypi

之前的方法是print(model)，或者在debug的时候去取值，但是在pip install torchinfo ，可以很方便的查看网络中的信息：

# 官方的例子
from torchinfo import summary
import torchvision

model = torchvision.models.resnet152()
summary(model, (1, 3, 224, 224), depth=3)

# [看看之前的vivit的效果](https://blog.csdn.net/ResumeProject/article/details/123470594?)
# 效果还行
	img = torch.ones([1, 16, 3, 64, 64]).cuda()
    # b t c (h p1) (w p2) -> b t (h w) (p1 p2 c)   p1=p2=16  # torch.Size([1, 16, 3, 64, 64]) -> torch.Size([1, 16, 16, 192])
    model = ViViT(224, 16, 100, 16).cuda()
    from torchinfo import *
    summary(
        model,                 # PyTorch model
        (1, 16, 3, 64, 64),    #  Shape of input data as a List/Tuple/torch.Size
        # dtypes=[torch.long],
        # verbose=2,
        # col_width=16,
        # col_names=["kernel_size", "output_size", "num_params", "mult_adds"],
        # row_settings=["var_names"],
    )

====================================================================================================
Layer (type:depth-idx)                             Output Shape              Param #
====================================================================================================
ViViT                                              --                        --
├─Transformer: 1                                   --                        --
│    └─ModuleList: 2-1                             --                        --
│    │    └─ModuleList: 3-1                        --                        444,288
│    │    └─ModuleList: 3-2                        --                        444,288
│    │    └─ModuleList: 3-3                        --                        444,288
│    │    └─ModuleList: 3-4                        --                        444,288
├─Transformer: 1                                   --                        --
│    └─ModuleList: 2-2                             --                        --
│    │    └─ModuleList: 3-5                        --                        444,288
│    │    └─ModuleList: 3-6                        --                        444,288
│    │    └─ModuleList: 3-7                        --                        444,288
│    │    └─ModuleList: 3-8                        --                        444,288
├─Sequential: 1-1                                  [1, 16, 16, 192]          --
│    └─Rearrange: 2-3                              [1, 16, 16, 768]          --
│    └─Linear: 2-4                                 [1, 16, 16, 192]          147,648
├─Dropout: 1-2                                     [1, 16, 17, 192]          --
├─Transformer: 1-3                                 [16, 17, 192]             --
│    └─LayerNorm: 2-5                              [16, 17, 192]             384
├─Transformer: 1-4                                 [1, 17, 192]              --
│    └─LayerNorm: 2-6                              [1, 17, 192]              384
├─Sequential: 1-5                                  [1, 100]                  --
│    └─LayerNorm: 2-7                              [1, 192]                  384
│    └─Linear: 2-8                                 [1, 100]                  19,300
====================================================================================================
Total params: 3,722,404
Trainable params: 3,722,404
Non-trainable params: 0
Total mult-adds (M): 30.39
====================================================================================================
Input size (MB): 0.79
Forward/backward pass size (MB): 20.37
Params size (MB): 14.89
Estimated Total Size (MB): 36.05
====================================================================================================

Process finished with exit code 0