7

pytorch模型结构可视化，可显示每层的尺寸

 3 years ago
source link: http://www.cnblogs.com/haoliuhust/p/14257868.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

最近在学习一些检测方面的网络，使用的是pytorch。模型结构可视化是学习网络的有用的部分，pytorch没有原生支持这个功能，需要找一些其他方式，下面总结几种方法 （推荐用4） 。

1. torch .pt -> netron

netron是一个专门可视化模型的工具，支持很多格式，很自然的就是用它直接显示torch保存的模型。但是实际上...

INbYVfn.png!mobile

显示成了上图，基本上没什么用。

2. 网上说的比较多的几种方式

a. tensorboardx 画出来的比较丑

b. tensorwatch 支持的torch版本低，容易出错

c. graphviz 实测yolov5异常。。

3. torch->caffe->netron

caffe支持的操作不够多，转换很容易碰到不支持的情况

4. torch->onnx->netron

pytorch对onnx支持的比较好，直接有导出接口。

torch.onnx.export...

但这样直接转出来的onnx进行显示没有每一层的输入大小，如下图，还是不方便看。

VbIjeuj.png!mobile

我们可以将导出的onnx进行以下操作：

import onnx

from onnx import shape_inference

model = 'best.onnx'

onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)

再用netron进行显示即可得到，用这种结构图学习就好多了：

6zeUjem.png!mobile

欢迎关注：老司机的视觉屋


report

Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK