飞桨x昇腾生态适配方案：14_loop算子缺失（上）：ONNX模型拆分

方案背景

当在线推理的速度无法满足客户要求，使用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"["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工具打开后模型结构如图所示：