方案背景
当在线推理的速度无法满足客户要求，使用 atc 工具将 onnx 转为 om 模型走离线推理路径时，遇到 NPU 不支持 LOOP 算子的问题，本文提供一种解决方案。本方案的设计思路是，onnx 文件分成 loop 算子和不含 loop 算子的两部分，把 loop 算子的子图提取出来，单独推理。实际操作中可能需要分成 3 份乃至更多，因此，本方案使用于关键路径上的 loop 算子，否则工作量会很大。
构造包含 loop 算子的 onnx 模型
首先使用以下代码构造一个包含 loop 算子的 onnx 模型
# 文件名：hi.pyimport torch
@torch.jit.scriptdef 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:
执行结束后生成 model.onnx 文件，使用 netron 工具打开后模型结构如图所示：
onnx 文件转换为 json 格式
# 文件名：onnx_to_json.pyimport onnxfrom google.protobuf.json_format import MessageToJsonmodel = onnx.load("./model.onnx") message = MessageToJson(model)with open("{}.json".format("./model.onnx"), "w") as fi: fi.write(message)
复制代码
执行以上脚本文件onnx_to_json.py:
执行结束后生成 model.onnx.json 文件，打开此文件可以看到，Loop 算子包含了一个子图，后续的工作是将子图提取出来
提取子图
替换主图
用 loop["attribute"][0]["g"] 的内容替换主图 "graph"中的内容
json 转换成 onnx
把替换之后 json 转换成 onnx，得到如下 onnx 文件：
# 文件名：jsonTo_onnx.pyimport onnximport numpy as npfrom google.protobuf.json_format import MessageToJson, Parseimport jsonimport time
# json to onnxwith 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:
执行结束后生成 model_loop.onnx 文件，使用 netron 工具打开后模型结构如图所示：
切分原图
原图分为：loop 算子前面的图-A、loop 算子-B、loop 算子后的图-C，这三部分，需要加载 3 个图，然后在 loop 算子那块循环做 infer 本例原图的这个 loop 在结尾的位置，那就不需要切分 loop 算子后的图-C，只需要切分成 loop 算子前面的图-A、loop 算子-B
# 文件名：extract_model.pyimport onnxfrom onnx.utils import extract_modelonnx.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 工具打开后模型结构如图所示：
评论