方案背景
当在线推理的速度无法满足客户要求，使用atc工具将onnx转为om模型走离线推理路径时，遇到NPU不支持LOOP算子的问题，本文提供一种解决方案。
本方案的设计思路是，onnx文件分成loop算子和不含loop算子的两部分，把loop算子的子图提取出来，单独推理。
实际操作中可能需要分成3份乃至更多，因此，本方案使用于关键路径上的loop算子，否则工作量会很大。
构造包含loop算子的onnx模型
首先使用以下代码构造一个包含loop算子的onnx模型
# 文件名：hi.py
import torch
@torch.jit.script
def loop(x, y):
for i in range(x.shape[0]):
x = x + y
return x
class SimpleModel(torch.nn.Module):
def _mul(self, x, y = 20):
return x * y
def forward(self, x, y):
x = x - y
x = loop(x, y)
return x
if __name__ == "__main__":
input1 = torch.rand((2, 16, 32))
input2 = torch.rand((2, 16, 32))
model = SimpleModel()
torch.onnx.export(
model,
(input1, input2),
'./model.onnx',
input_names=['input1', 'input2'],
dynamic_axes = {
'input1': {0: 'batch'},
'input2': {0: 'batch'},
}
)
执行以上脚本文件hi.py:
python3 hi.py
执行结束后生成model.onnx文件，使用netron工具打开后模型结构如图所示：
onnx文件转换为json格式
# 文件名：onnx_to_json.py
import onnx
from google.protobuf.json_format import MessageToJson
model = onnx.load("./model.onnx")
message = MessageToJson(model)
with open("{}.json".format("./model.onnx"), "w") as fi:
fi.write(message)
执行以上脚本文件onnx_to_json.py:
python3 onnx_to_json.py
执行结束后生成model.onnx.json文件，打开此文件可以看到，Loop算子包含了一个子图，后续的工作是将子图提取出来
提取子图
替换主图
用 loop["attribute"][0]["g"] 的内容替换主图 "graph"中的内容
json转换成onnx
把替换之后json转换成onnx，得到如下onnx文件：
# 文件名：jsonTo_onnx.py
import onnx
import numpy as np
from google.protobuf.json_format import MessageToJson, Parse
import json
import time
# json to onnx
with open("model_loop.onnx.json", "r") as fi:
onnx_json = json.loads(fi.read())
onnx_str = json.dumps(onnx_json)
model2 = Parse(onnx_str, onnx.ModelProto())
onnx.save(model2, "model_loop.onnx")
执行以上脚本文件jsonTo_onnx.py:
python3 jsonTo_onnx.py
执行结束后生成model_loop.onnx文件，使用netron工具打开后模型结构如图所示：
切分原图
原图分为：loop算子前面的图-A、loop算子-B、loop算子后的图-C，这三部分，需要加载3个图，然后在loop算子那块循环做infer
本例原图的这个loop在结尾的位置，那就不需要切分loop算子后的图-C，只需要切分成loop算子前面的图-A、loop算子-B
# 文件名：extract_model.py
import onnx
from onnx.utils import extract_model
onnx.utils.extract_model("model.onnx","model_dest1.onnx",['input1', 'input2'], ['/Sub_output_0', 'onnx::Loop_6'], check_model=False)
['input1', 'input2']:切分模型的输入节点
['/Sub_output_0', 'onnx::Loop_6']：切分模型的输出节点
执行以上脚本文件extract_model.py，生成loop算子前面的图-A:
修改json构造子图
由于Add算子的另一个输入直接来自主图的input2，但我们直接提取的子图中没有input2，故需要添加一个input2节点
在子图的json文件中添加来自主图的input2：
json转onnx文件
执行脚本文件jsonTo_onnx.py，生成新带有input2的子图：
构造子图输入
- 输入 i，在实际在循环控制中发挥作用，此处不生效，任意传入一个值占位即可，可删除节点
- 输入cond
- inputs = {"i": 0 , cond: True, "x.13": outputs[0], "input2":input2}
loop的执行逻辑
// 当迭代次数小于最大行程计数，并且条件的 ML 值所指向的布尔张量数据为 true 时，进入循环
while (iter_num_value < max_trip_count_ && *condition_mlvalue_.GetMutable<Tensor>()->MutableData<bool>()) {
// 如果迭代次数不为 0
if (iter_num_value != 0) {
// 保存输出并更新输入
SaveOutputsAndUpdateFeeds(fetches, feeds);
// 清空 fetches 向量
fetches.clear();
}
// 执行子图，将执行结果存储在 status 中
status = utils::ExecuteSubgraph(session_state_, ffm, feeds, fetches, {},
ExecutionMode::ORT_SEQUENTIAL, context_.GetTerminateFlag(), context_.Logger(),
context_.GetComputeStream(),
// 由于 fetches[0] 是循环条件，我们需要在 CPU 上访问该数据，
// 因此必须进行流同步以确保数据已到达。
true);
// 如果执行子图过程中出现错误，则返回错误
ORT_RETURN_IF_ERROR(status);
// 将 fetches 向量中的第一个元素赋值给 condition_mlvalue_
condition_mlvalue_ = fetches[0];
以上代码实现了一个循环，在满足特定条件时会持续执行子图。每次迭代时，若不是第一次迭代，就会保存输出并更新输入，接着执行子图，最后更新循环条件。
删除无关节点
循环控制分为循环次数和循环condition，子图的第一个输出为condition。本例中condition始终为True，可以删除这个节点。
在子图json文件中删除图中圈出的输入数据节点：
删除后执行脚本文件jsonTo_onnx.py，执行结束后生成mode_loop_input2_i_cond.onnx文件，使用netron工具打开后模型结构如图所示：
