Pytorch神经网络结构可视化模块--Tensorwatch

Tensorwatch

可以用于pytorch可视化的模块还有netron、tensorboardx、pytorchviz。
netron对pytorch模型的支持不友好，pytorchviz显示的效果又不佳

Tenosrwatch实现起来很简单
Tensorwatch在jupyter notebook中可以直接显示结构图片，在idle中需要手动保存，手动保存的代码在后面，推荐保存到本地查看。

Tensorwatch的项目地址：https://github.com/microsoft/tensorwatch

环境

Win10 + Anaconda
python = 3.7.3
pydot =1.4.1
graphviz = 2.38
pytorch = 1.2
scikit-image = 0.15.0
scikit-learn = 0.21.2
scipy = 1.2.1
tensorwatch = 0.8.7
torchvision = 0.4.0

安装

pip install pydot
pip install tensorwatch

错误一

另外需要注意如果报了：AttributeError: ‘Dot’ object has no attribute ‘repr_svg’ 的错误
将 Anaconda\Lib\site-packages\tensorwatch\model_graph\hiddenlayer\pytorch_draw_model.py的
第13行改为 return self.dot.create_svg().decode()

错误二

如果继续报：FileNotFoundError: [WinError 2] “dot” not found in path.
用conda安装graphviz

conda isntall graphviz

使用

tensorwatch支持查看网络的结构和网络的统计结果

查看网络结构

import tensorwatch as tw
import torchvision.models
alexnet_model = torchvision.models.alexnet()
tw.draw_model(alexnet_model, [1, 3, 224, 224])

保存网络结构图

import tensorwatch as tw
import torchvision.models
alexnet_model = torchvision.models.alexnet()
img = tw.draw_model(alexnet_model, [1, 3, 1024, 1024])
img.save(r'C:/Users/Hassaki/Desktop/alexnet.jpg')

查看网络的统计结果

import tensorwatch as tw
import torchvision.models
alexnet_model = torchvision.models.alexnet()
tw.model_stats(alexnet_model, [1, 3, 224, 224])