接上一章节内容，将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：
可看到推理结果一致。
