ONNX节点增、删、改、查（网络结构修改）

onnx网络结构修改，便于将模型部署到一些不支持某些特定op的边缘（嵌入式）设备

---

前言

onnx网络结构修建的目的

当我们需要把模型部署到一些嵌入式设备时，往往需要采用ncnn、OpenVINO、PPLNN、RKNN 、SNPE、
TensorRT、TVM 、Core ML 等部署方案。而这些部署方案往往需要先将Pytorch或TensorFlow模型转成onnx模型，进而再转换成对应的部署框架支持的模型格式。而由于硬件设备、厂商支持力度等原因的存在，上述多种部署框架大多存在一些不支持的op（算子），导致部署模型之路困难重重，因此，在onnx模型是进行适当的“手术”，对不支持的op进行增删改查是十分必要的。
本文将重点叙述如何对onnx模型进行“手术”，以满足部署端的需求。

---

一、python的onnx_graphsurgeon 库？

onnx_graphsurgeon 是模型部署工程的必备工具之一, 最主要的用途就是调整或者新建 onnx 模型
ONNX GraphSurgeon is a tool that allows you to easily generate new ONNX graphs, or modify existing ones.

具体的工作场景比如, 模型算法小组训练好了一个模型, 然后把这个模型交到了新来的同事手里让她部署测试一下, 她接到模型打开一看, 好家伙, onnx版本不兼容, 或者打开模型一看, 好家伙, 算子属性配置错了. 正在她不知所措之际, 你拍了拍她稚嫩的肩膀, 将 onnx_graphsurgeon 递了过去. 她无助的眼神突然闪烁起了希望。

二、使用步骤

1.安装必要的库

安装 onnx_graphsurgeon命令如下：

$ python3 -m pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

2.使用场景（删除rtmpose模型中的Abs算子）

原始网络结构（有Abs算子，部署框架不支持，因此需要更改）：

修改之后的网络结构（删除了Abs算子，可以部署）：

当然，我们也可以选择更改Aba算子为Constant算子或其他合适的算子：

1.所有代码

代码如下：

import onnx_graphsurgeon as gs
import onnx
    def remove_abs(model_path):


        graph = gs.import_onnx(onnx.load(model_path))

        remove_node_name = '/mlp/mlp.0/Abs'
        remove_node = [
            node for node in graph.nodes if node.name == remove_node_name
        ][0]
        inp_node = remove_node.i()
        oup_node = remove_node.o()
        oup_node.inputs[0] = inp_node.outputs[0]
        remove_node.outputs.clear()

        # Remove the fake node from the graph completely
        graph.cleanup()
        onnx.save(gs.export_onnx(graph), model_path)
        onnx_model = onnx.load(model_path)
        onnx.checker.check_model(onnx_model)


    remove_abs(onnx_model)

---

总结

如需任何算法开发、部署、模型转换等帮助，可联系我。
微信：One-of-immortals
QQ：584746814