三种实现模型可视化的方式(print, torchinfo, tensorboard)

记录一下自己使用的三种模型可视化的方式，从简单到难

Print

最简单的是print，就不用多说了。

Torchinfo

from torchinfo import summary
import torch
model = (...)
summary(model, (1,3,128,128))

即可按照像文档路径一样的方式输出结构，并且有每一层的shape和参数量，比print的信息更多。

但相对而言还是比较简单。

Tensorboard

该方法是个人觉得最好用的并且很动态
需要提前下好tensorboard和tensorboardX

import torch
from torchvision.models import AlexNet
 
from tensorboardX import SummaryWriter
 
 
x=torch.rand(1,3,256,512)
model=AlexNet()
 
with SummaryWriter(comment='AlexNet') as w:
    w.add_graph(model, x) 

运行该方法后会在同级路径下生成一个runs文件夹，其中就有每次的文件

再开一个终端，进入runs之后，输入指令tensorboard --logdir=XXX --port 8080 这里XXX就是runs下方某一个文件夹的名称。--port 指定端口。个人建议指定一个不常用的端口，因为有时常用端口会被其他程序占用，不一定能显示，遇到过类似的坑。运行了命令之后点开对应的链接，一般是localhost:port 或者127.0.0.1:port ，浏览器打开即可。

他这个图是动态的，点击你想看的模块可以放大，比如打开Alexnet，只要有模块就可以一直方法，很方便。

相信这三种方法已经够你们使用了！