如何在OpenVINO中實現自己的自定義運算元 – 第一篇：搞定模型轉換
前言
與主流的深度學習框架類似，OpenVINO也提供一套擴充套件機制，方便開發者可以實現自己的自定義運算元。自定義運算元有如下三種使用場景：
- 模型中含有MO不支援的運算元，但該運算元可以通過已支援運算元的組合實現，這時自定義運算元實際上完成的是將“不支援的運算元”替換為“支援運算元”的實現。
- 模型中含有的子圖可以被替換為效能更佳的其它等價運算，這種操作被OpenVINO稱為”Fuse”(融合)。舉例來講，比如Swish操作，就會被等價替換為計算式：
- 模型中含有原生框架的自定義運算元，比如本文中就會利用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". [ ERROR ] Op type not registered 'ZeroOut' in binary running on test-ubuntu. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed. [ ERROR ] [ ERROR ] It can happen due to bug in custom shape infer function <function tf_native_tf_node_infer at 0x7f4515c95378>. [ ERROR ] Or because the node inputs have incorrect values/shapes. [ ERROR ] Or because input shapes are incorrect (embedded to the model or passed via --input_shape). [ ERROR ] Run Model Optimizer with --log_level=DEBUG for more information. [ ERROR ] Exception occurred during running replacer "REPLACEMENT_ID" (<class 'extensions.middle.PartialInfer.PartialInfer'>): Stopped shape/value propagation at "ModZeroOut/Zeroout/zero_out" node. For more information please refer to Model Optimizer FAQ, question #38. (https://docs.openvinotoolkit.org/latest/openvino_docs_MO_DG_prepare_model_Model_Optimizer_FAQ.html?question=38#question-38)
提示說”ZeroOut”沒有被註冊。接下來我們修復這個問題。
認識工作目錄
首先認識一下mo擴充套件程式碼的工作目錄：
cd /opt/intel/openvino_2021.2.185/deployment_tools/model_optimizer/extensions其中包括如下子目錄：
其中:
analysis是模型的分析程式碼，包含了對yolo, fasterrcnn等的預分析程式碼，感興趣的同學可以自己深入研究。
load, front, middle, back分別對應於MO轉換的四個階段，官網上可以找到對應的階段說明，這裡講一些我自己的理解：
load: 對應於官網中的Model Loading階段，是讀取原生框架模型的模組，裡面包含有tf, onnx, mxnet, caffe, kaldi相對應的實現，這些屬於基礎程式碼，除非我們想支援新的框架，否則不需要去修改這裡。
front: 對應於官網中的Front Phase，這裡是MO轉換的最早期階段，可以對運算元的輸入進行處理，比如extensions/front/TopKNormalize.py裡的處理，是將TopK中的屬性k移走，並新增一個常量k的輸入節點。實際上， Font Phase裡可以做非常多的處理，比如對op進行替換，也可以放在這個階段進行。
middle: 對應於官網中的Middle Phase，一般在font phase裡，經過一系列的變換，可能會有新產生的op，在middle phase，需要更新所有op的屬性，或者是合併/刪除一些冗餘的操作。
back: 對應於官網中的Back Phase, 一般來講在middle phase裡處理不了的情況，就可以放在back phase裡來處理。
ops: 放的是運算元的基類，比如extensions/ops/topk.py，我這裡要實現的ZeroOut運算元就放在這個目錄裡。
其實對於這裡的子目錄，更多的我認為是需要遵循的編碼規範，比如說如果我想新增加一個複雜op的支援，並且希望可以合併到openvino的master裡，那麼不同的程式碼需要放在不同的目錄下。
MO中的具體實現
我這裡需要實現的功能非常簡單，只是希望MO可以認識ZeroOut這個op，從面成功生成IR檔案，不需要對ZeroOut進行任何的特殊處理，所以我只在ops裡增加了ZeroOut的支援。
下面是extensions/ops/ZeroOut.py的簡單實現：
from mo.front.common.partial_infer.elemental import copy_shape_infer
from mo.graph.graph import Graph
from mo.ops.op import Op
class ZeroOutOp(Op):
op = 'ZeroOut'
def __init__(self, graph: Graph, attrs: dict):
mandatory_props = {
'type': __class__.op,
'op': __class__.op,
'version': 'extension',
'infer': copy_shape_infer
}
super().__init__(graph, mandatory_props, attrs)
添加了ZeroOut的實現後，再次執行轉換命令：
mo.py -m frozen_graph.pb -o ir/FP32/MO轉換成功：
Model Optimizer version: 2021.2.0-1877-176bdf51370-releases/2021/2
[ SUCCESS ] Generated IR version 10 model.
[ SUCCESS ] XML file: /home/test/Workspace/VMC/OpenVINO/r2021.1/custom-layer/create_tf_model/tf_model/ir/FP32/frozen_inference_graph.xml
[ SUCCESS ] BIN file: /home/test/Workspace/VMC/OpenVINO/r2021.1/custom-layer/create_tf_model/tf_model/ir/FP32/frozen_inference_graph.bin
[ SUCCESS ] Total execution time: 3.37 seconds.
[ SUCCESS ] Memory consumed: 399 MB.
為了驗證確實轉換成功了，我們開啟xml檔案看一下，能找到如下的ZeroOut運算元，版本是extension：
<layer id="61" name="ModZeroOut/Zeroout/zero_out" type="ZeroOut" version="extension">
<input>
<port id="0">
<dim>1</dim>
<dim>2</dim>
</port>
</input>
<output>
<port id="1" precision="FP32">
<dim>1</dim>
<dim>2</dim>
</port>
</output>
</layer>
可以看到ZeroOut這個運算元已經被MO所識別，並且生成了ir。
後記
總的來說，MO起到的最主要作用是生成帶有運算元屬性的xml檔案。其實還有進階用法，在官網上有一個更復雜的例子：https://docs.openvinotoolkit.org/latest/openvino_docs_HOWTO_Custom_Layers_Guide.html，大家感興趣的話可以再研究一下。後續我也會找時間再詳細解讀這個複雜的例子。
不過，生成ir只是支援自定義運算元的第一步，因為在推理階段，推理引擎還需要識別這個運算元，接下來我會講到如何在intel CPU和 intel整合顯示卡上實現自定義運算元的推理。
參考連結：
https://docs.openvinotoolkit.org/latest/openvino_docs_IE_DG_Extensibility_DG_Intro.html
https://github.com/openvinotoolkit/openvino/tree/master/docs/template_extension
https://docs.openvinotoolkit.org/latest/openvino_docs_HOWTO_Custom_Layers_Guide.html