大家好,我是极智视界,本文介绍一下 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模型增、删、改、查算子节点方法。希望我的分享能对你的学习有一点帮助。