网络结构可视化onnx+netron

ResNet18网络结构可视化效果：

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为深度学习模型可视化库，支持多种模型格式，包括ONNX等。将onnx模型拖入Netron官网即可查看可视化结果。

教程：netron 神经网络可视化教程 - 知乎

网页使用界面：Netron

onnx-simplifier

将pytorch等框架的网络模型转化为onnx模型进行可视化时，会存在冗余现象，因此需要onnx-simplifier来减少冗余，进一步提高可视化结构可读性。

如下图，对输入shape(2,3,4,5)进行reshape操作变为(2,3,5,4)的冗余可视化结果：

减少冗余后：

说明：ONNX学习笔记 - 知乎

github： https://github.com/daquexian/onnx-simplifier#python-version

安装：

pip install -U pip && pip install onnx-simplifier

使用：

onnxsim  

代码

对上述库进行了简单的封装，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')