summary打印模型结构(VIT为例)

问题：怎么看部署的ViT结构？
解决：实例化生成的模型，结合summary方法打印，以pytorch环境为例：
在虚拟环境中下载对应的torchsummary，pip安装指令：pip install torchsummary。
测试是否成功：输入‘python’,进入IDLE，import torchsummary不报错则成功。

打印VIT结构：

import torch
from torchsummary import summary

from vit_model import vit_base_patch16_224_in21k as create_model

if __name__ == "__main__":
    # 需要使用device来指定网络在GPU还是CPU运行
    device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    #m       = YoloBody([[6, 7, 8], [3, 4, 5], [0, 1, 2]], 80).to(device)
    model = create_model(num_classes=5, has_logits=False).to(device)
    summary(model, input_size=(3, 224, 224))

"""
#打印yolo方法
import torch
from torchsummary import summary

from nets.yolo import YoloBody

if __name__ == "__main__":
    # 需要使用device来指定网络在GPU还是CPU运行
    device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    m       = YoloBody([[6, 7, 8], [3, 4, 5], [0, 1, 2]], 80).to(device)
    summary(m, input_size=(3, 416, 416))
"""

部分结果：

其实summary还有很多实用的用法：
【相关资源】Tensorflow学习笔记——Summary用法