飞桨x昇腾生态适配方案：15_loop算子缺失（下）：for循环替换loop

接上一章节内容，将ONNX模型拆分成loop算子部分和非loop算子部分后，分别转换成OM模型，并用for循环替换loop算子计算逻辑，比较OM模型和ONNX模型的推理结果是否一致，验证结果如果一致则证明该方案有效。

onnx模型转om

loop算子前面的图-A

atc --model=./mode_loop_input2_i_cond.onnx  --framework=5 \
    --output=./mode_loop_input2_i_cond --soc_version=Ascend910B2 \
    --input_shape="input1:1~8,16~32,16~32;input2:1~8,16~32,16~32" \
    --input_format=ND --log=error

loop算子子图-B

atc --model=./mode_loop_input2_i_cond.onnx  --framework=5 \
    --output=./mode_loop_input2_i_cond --soc_version=Ascend910B2 \
    --input_shape="x.13:1~8,16~32,16~32;input2:1~8,16~32,16~32" \
    --input_format=ND --log=error

构造子图执行逻辑

for循环实现loop计算逻辑

output1=before_loop.infer(input)
for 循环：
    output2=子图执行(output1（首次）/output2（后续）)
output3=after_loop.infer(outout2)

编写om模型推理脚本

根据以上逻辑编写本例for循环：

# 文件名：loop.py
import numpy as np
from ais_bench.infer.interface import InferSession

# Load the model
model_path0 = './model_dest_linux_x86_64.om'
model_path1 = './mode_loop_input2_i_cond_linux_x86_64.om'

# input
device_id = 0
session0 = InferSession(device_id, model_path0)
session1 = InferSession(device_id, model_path1)
# 假设 batch 大小为 4
batch = 4

# 生成 input1 和 input2
# input1 = np.random.randn(batch, 16, 32).astype(np.float32)
# input2 = np.random.randn(batch, 16, 32).astype(np.float32)

# print("input1 的形状:", input1.shape)
# print("input2 的形状:", input2.shape)

# 生成全 1 的 input1
input1 = np.ones((batch, 16, 32), dtype=np.float32)

# 生成全 2 的 input2
input2 = np.full((batch, 16, 32), 2, dtype=np.float32)

print("input1 的形状:", input1.shape)
print("input2 的形状:", input2.shape)
print("input1 的前几个元素:", input1[0, 0, :5])
print("input2 的前几个元素:", input2[0, 0, :5])

input = [input1, input2]
#
output0 = session0.infer(input, mode='dymshape')
input3 = output0[0]
cnts = output0[1]
# print("input3", input3)
# print("cnts", cnts)

for i in range(cnts):
    inputs = [input3, input2]
    outputs = session1.infer(inputs,mode='dymshape')
    input3 = outputs[0]
    # print(input3)
print("outputs", outputs)

执行python3 loop.py:

验证拆分om与原图onnx结果

编写onnx模型推理脚本

# 文件名：tensor_generation.py
import numpy as np
import onnxruntime as ort

# 假设 batch 大小为 4
batch = 4

# 生成随机的 input1 和 input2
input1 = np.random.randn(batch, 16, 32).astype(np.float32)
input2 = np.random.randn(batch, 16, 32).astype(np.float32)

# 加载 ONNX 模型
model_path = './model.onnx'  # 替换为你的 ONNX 模型文件路径
session = ort.InferenceSession(model_path)

# 获取输入名称
input_names = [input.name for input in session.get_inputs()]

# 准备输入字典
if len(input_names) == 2:
    inputs = {input_names[0]: input1, input_names[1]: input2}
elif len(input_names) == 1:
    # 如果模型只有一个输入，这里假设使用 input1
    inputs = {input_names[0]: input1}
else:
    raise ValueError(f"模型输入数量为 {len(input_names)}，不支持，请检查模型。")

# 进行推理
outputs = session.run(None, inputs)

# 输出结果
print("推理结果：", outputs)
for i, output in enumerate(outputs):
    print(f"输出 {i} 的形状: {output.shape}")

执行tensor_generation.py：

可看到推理结果一致。