如何在OpenVINO中实现自己的自定义算子 – 第一篇：搞定模型转换

本文链接：https://blog.youkuaiyun.com/cosmos_qiu/article/details/113756176

本文介绍了如何在OpenVINO中实现自定义算子，以解决模型转换时遇到的不支持算子问题。通过ZeroOut算子为例，详细阐述了从模型转换的错误提示到理解工作目录，再到具体实现ZeroOut算子的Python代码，并最终成功生成IR文件的过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

与主流的深度学习框架类似，OpenVINO也提供一套扩展机制，方便开发者可以实现自己的自定义算子。自定义算子有如下三种使用场景：

模型中含有MO不支持的算子，但该算子可以通过已支持算子的组合实现，这时自定义算子实际上完成的是将“不支持的算子”替换为“支持算子”的实现。
模型中含有的子图可以被替换为性能更佳的其它等价运算，这种操作被OpenVINO称为”Fuse”(融合)。举例来讲，比如Swish操作，就会被等价替换为计算式： $x/(1.0+e^{^{-(beta*x)}})$
模型中含有原生框架的自定义算子，比如本文中就会利用tf实现一个叫ZeroOut的自定义算子，这个算子在tf里作了自定义实现。同时，在MO里就需要定义一个同样的ZeroOut自定义算子，从而让MO可以认识ZeroOut算子。

对普通开发者而言，第3种情况是最常遇见的 – 因为原生框架不支持而不得不自己实现 – 实际上不管哪种情况实现原理都是一样的，下面就讲讲我自己是怎么实现的。

准备工作

我使用的环境为：Ubuntu18.04, OpenVINO2021.2, tf1.15.2

在开始之前，首先需要准备好一个含有ZeroOut算子的网络模型，这里我实现了一个空的网络并进行了冻结，pb文件可以从这里下载：https://github.com/cosmos1982/OpenVINO_Demos/blob/master/custom-layer/create_tf_model/tf_model/frozen_inference_graph.pb

ZeroOut的数学公式和实现，参考tf官网：https://www.tensorflow.org/guide/create_op?hl=zh-cn

在什么也不修改的情况下，我使用MO进行模型转换，转换命令：

mo.py -m frozen_graph.pb -o ir/FP32/

会得到如下错误：

Model Optimizer version:        2021.2.0-1877-176bdf51370-releases/2021/2
[ ERROR ]  Cannot infer shapes or values for node "ModZeroOut/Zeroout/zero_out"