飞桨 x 昇腾生态适配方案：15_loop 算子缺失（下）：for 循环替换 loop
作者：小顺637
- 2025-05-12 北京
本文字数：1601 字
阅读完需：约 5 分钟
接上一章节内容，将 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.pyimport numpy as npfrom ais_bench.infer.interface import InferSession
# Load the modelmodel_path0 = './model_dest_linux_x86_64.om'model_path1 = './mode_loop_input2_i_cond_linux_x86_64.om'
# inputdevice_id = 0session0 = InferSession(device_id, model_path0)session1 = InferSession(device_id, model_path1)# 假设 batch 大小为 4batch = 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 的 input1input1 = np.ones((batch, 16, 32), dtype=np.float32)
# 生成全 2 的 input2input2 = 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.pyimport numpy as npimport onnxruntime as ort
# 假设 batch 大小为 4batch = 4
# 生成随机的 input1 和 input2input1 = 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：
可看到推理结果一致。
划线
评论
复制
发布于: 刚刚阅读数: 4
小顺637
关注
还未添加个人签名 2023-01-19 加入
还未添加个人简介
评论