本文总结了自己关于PyTorch模型结构可视化的一些尝试,如若描述存在纰漏,欢迎各位指正交流。
PyTorch模型结构可视化,几种方法按照推荐顺序排列如下:
1. netron(个人最推荐)
- 具体方法:pytorch存下来的.pth文件转为.onnx文件,直接拖入netron网页就可以看
- 优点:画出来比较细致,布局效果好,直观
2. tensorboard(个人不太推荐)
- 具体方法:安装合适版本的PyTorch\tensorboard,在代码中添加summary、add_graph
- 缺点:
a. 加上tensorboard中关于summarywriter的代码后,好像会让显存增加不少,而且第一张卡的显存特别多
b. 需要PyTorch较高版本
c.画图效果一般(可能会因为版本原因,出现画图错误的问题(画出来只有一个空白内容的框),或者画图)