极智AI | onnx模型增删改查算子节点方法

Posted 2022-12-16 极智视界

  欢迎关注我，获取我的更多笔记分享

  大家好，我是极智视界，本文介绍一下 onnx 模型增、删、改、查算子节点方法。

  大家知道，onnx 是 AI 业界用的很多的 中间模型 (应该没人直接把 onnx 作为原始模型用吧)，一般的初始训练模型为 pytorch 或 tensorflow，然后再转换为 onnx ，然后 再用作他用，比如拿来部署等。我画了一个示意图：

  这个图看着可能会感觉有些疑惑，我来解释一下：首先左边是训练框架，右边是推理框架，总的 flow 应该是 左->右，但我又画了几个回旋镖，这些回旋镖的意思是：训练框架模型之间也有可能需要通过 onnx 进行转换。

  当然这里不是主要讲 模型之间的转换，而是讲 有了onnx模型后，怎么对其中的算子节点进行一些自定的操作呢？

  下面开始。

  假设我们已经有了一个 resnet50.onnx (如果没有，可以很方便的从 pytorch modelzoo 中导入，然后通过 torch.onnx.export 导出)。

import onnx

# 导入resnet50.onnx模型
resnet50_onnx = onnx.load("./resnet50.onnx")
# 获得onnx图
graph = resnet50_onnx.graph
# 获得onnx节点
node = graph.node

### 准备工作已就绪，开干
# 增、删、改、查一起操作
# 比如咱们要对 `算子类型为Add&输出为225的节点` 进行操作
for i in range(len(node)):
    if node[i].op_type == 'Add':
        node_rise = node[i]
        if node_rise.output[0] == '225':
            print(i)  # 169 => 查到这个算子的ID为169

old_node = node[169]  # 定位到刚才检索到的算子

# 新增一个 `Constant` 算子
new_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=['225'],
    value=onnx.helper.make_tensor('value', onnx.TensorProto.FLOAT, [4], [1, 1, 1.2, 1.2])
)  

# 删除旧节点
node.remove(old_node)  

# 插入新节点
node.insert(169, new_node)  

# 是不是还少一个修改节点，比方看下面
# node[169].type = 'Conv'   # 将刚才的算子类型改为2D卷积
# 改名称啥的类似

### 保存新模型
# 校验
onnx.checker.check_model(resnet50_onnx)
# 保存
onnx.save(resnet50_onnx,'resnet50_new.onnx')

  说一下 onnx.checker.check_model，这个接口是用来校验新生成的 onnx 模型的 graph、node是否正确，比如你修改的算子类型并不在 onnx 的算子库中，则这个接口会抛出类似 onnx.onnx_cpp2py_export.checker.ValidationError: No Op registered for xxx_opname with domain_version of 11 的报错提醒。这个时候你可能有三个选择：(1) 修改新增的自定义算子类型，使其在onnx的算子库中；(1) 在onnx算子库中新增实现，以覆盖新增的自定义算子；(3) 将错就错，不校验直接保存，自定义算子实现在下一框架/阶段进行。

  好了，以上分享了 onnx模型增、删、改、查算子节点方法。希望我的分享能对你的学习有一点帮助。

 【极智视界】

《极智AI | onnx 模型增、删、改、查算子节点方法)》

搜索关注我的微信公众号【极智视界】，获取我的更多经验分享，让我们用极致+极客的心态来迎接AI ！