效果展示
ResNet18网络结构可视化效果：
ONNX
Open Neural Network Exchange (ONNX) 可以将不同框架(TensorFlow/Pytorch/Paddle)训练出来的模型都转换onnx模型进行存储，起到统一模型格式的作用。
教程：ONNX学习 - 知乎
github：https://github.com/onnx/onnx
安装
满足：
numpy >= 1.16.6
protobuf >= 3.12.2
typing-extensions >= 3.6.2.1运行：
pip install numpy protobuf==3.16.0
pip install onnx
Netron
Netron为深度学习模型可视化库，支持多种模型格式，包括ONNX等。将onnx模型拖入Netron官网即可查看可视化结果。
网页使用界面：Netron
onnx-simplifier
将pytorch等框架的网络模型转化为onnx模型进行可视化时，会存在冗余现象，因此需要onnx-simplifier来减少冗余，进一步提高可视化结构可读性。
如下图，对输入shape(2,3,4,5)进行reshape操作变为(2,3,5,4)的冗余可视化结果：
减少冗余后：
github：?https://github.com/daquexian/onnx-simplifier#python-version
安装：
pip install -U pip && pip install onnx-simplifier使用：
onnxsim <input_onnx_model> <output_onnx_model>代码
对上述库进行了简单的封装，save_onnx.py：
# -*- coding:utf-8 -*-
"""
备注：将网络结构保存为onnx格式，便于netron可视化查看。
作者：lakuite
日期：2022年05月04日
"""
import torch
import onnx
from onnxsim import simplify
class Trans2ONNX():
def __init__(self, data, net, sim=True):
'''
:param data: 样例数据，torch.randn(b, c, h, w)
:param net: 网络，NET()
:param sim: 是否输出简化的可视化图，默认True
'''
super(Trans2ONNX, self, ).__init__()
self.data = data
self.net = net
self.sim = sim
def __call__(self, path):
'''
:param path: onnx输出路径，xxx.onnx
'''
# 导出为onnx格式
torch.onnx.export(self.net, self.data, path, export_params=True, opset_version=8)
model = onnx.load(path)
# 简化模型可视化
if self.sim:
model_simp, check = simplify(model)
assert check, "Simplified ONNX model could not be validated"
else:
model_simp = model
# 增加维度信息与保存
onnx.save(onnx.shape_inference.infer_shapes(model_simp), path)
调用：
import torchvision.models as models
import torch
from save_onnx import Trans2ONNX
# 定义样例数据+网络
data = torch.randn(2, 3, 224, 224)
net = models.resnet18()
ONNX = Trans2ONNX(data, net)
ONNX('onnx_model/resnet18.onnx')
