当前位置：首页 > news >正文

rknn转换后精度差异很大，失真算子自纠

news 来源：原创 2024/7/3 2:24:34

下面是添加了详细注释的优化代码：

import cv2
import numpy as np
import onnx
import onnxruntime as rt
from onnx import helper, shape_inferencedef get_all_node_names(model):"""获取模型中所有节点的名称。参数:model (onnx.ModelProto): ONNX 模型。返回:list: 包含所有节点名称的列表。"""return [node.name for node in model.graph.node]def remove_node_and_following(model, node_name):"""删除指定节点及其后续节点，并返回新的模型。参数:model (onnx.ModelProto): 原始 ONNX 模型。node_name (str): 要删除的节点名称。返回:onnx.ModelProto: 修改后的 ONNX 模型。"""nodes_to_keep = []  # 要保留的节点nodes_to_remove = set(i.name for i in model.graph.output)  # 要删除的节点start_removal = False  # 是否开始删除节点output = []  # 输出节点列表for node in model.graph.node:if node.name == node_name:start_removal = Trueif start_removal:nodes_to_remove.add(node.name)else:nodes_to_keep.append(node)output.extend(node.output)for node in model.graph.value_info:if node.name in output:shape = [dim.dim_value if (dim.dim_value > 0 and dim.HasField('dim_value')) else Nonefor dim in node.type.tensor_type.shape.dim]output_tensor = helper.make_tensor_value_info(node.name,onnx.TensorProto.FLOAT,shape)model.graph.output.append(output_tensor)new_graph = helper.make_graph(nodes_to_keep,model.graph.name,model.graph.input,[output for output in model.graph.output if output.name not in nodes_to_remove],model.graph.initializer,)new_model = helper.make_model(new_graph, producer_name=model.producer_name)new_model = shape_inference.infer_shapes(new_model)return new_modeldef preprocess_image(image_path, target_shape):"""加载并预处理图像。参数:image_path (str): 图像文件路径。target_shape (tuple): 目标形状 (宽, 高)。返回:np.ndarray: 预处理后的图像数组。"""im = cv2.imread(image_path)im = cv2.resize(im, target_shape)im = im.transpose((2, 0, 1))[::-1]  # HWC 转 CHW, BGR 转 RGBreturn np.ascontiguousarray(im)def main():model_path = 'yolov5s.onnx'model = onnx.load(model_path)dtype_map = {'tensor(float)': np.float32,'tensor(int32)': np.int32,'tensor(int64)': np.int64,}all_node_names = get_all_node_names(model)im = preprocess_image('./sample.jpg', (640, 640))for node_name in all_node_names:if node_name:new_model = remove_node_and_following(model, node_name)new_model_path = 'test.onnx'onnx.save(new_model, new_model_path)sess = rt.InferenceSession(new_model_path)output_names = [output.name for output in sess.get_outputs()]input_data = {input.name: im.reshape(input.shape).astype(dtype_map[input.type])for input in sess.get_inputs()}outputs = sess.run(output_names, input_data)# TODO: 补充 rknn 转换、加载和获取输出的代码# TODO: 补充评估对比两个输出 onnx 的 target 和 rknn 的 predict# TODO: 补充当评估准确率低于 0.8 时终止循环，此时的模型最后一个算子就是 rknn 不正确的算子if __name__ == "__main__":main()

在这里插入图片描述

解释

在补充 TODO 部分的代码之前，我们需要确保你有正确的 RKNN API 环境，并且你知道如何使用它来转换、加载和获取输出。以下是一个示例代码，展示了如何使用 RKNN API 完成这些任务。

1. 安装 RKNN Toolkit

首先，确保你已经安装了 RKNN Toolkit。你可以通过以下命令安装：

pip install rknn-toolkit

2. 补充代码

以下是补充了 TODO 部分的代码：

import cv2
import numpy as np
import onnx
import onnxruntime as rt
from onnx import helper, shape_inference
from rknn.api import RKNNdef get_all_node_names(model):"""获取模型中所有节点的名称。参数:model (onnx.ModelProto): ONNX 模型。返回:list: 包含所有节点名称的列表。"""return [node.name for node in model.graph.node]def remove_node_and_following(model, node_name):"""删除指定节点及其后续节点，并返回新的模型。参数:model (onnx.ModelProto): 原始 ONNX 模型。node_name (str): 要删除的节点名称。返回:onnx.ModelProto: 修改后的 ONNX 模型。"""nodes_to_keep = []  # 要保留的节点nodes_to_remove = set(i.name for i in model.graph.output)  # 要删除的节点start_removal = False  # 是否开始删除节点output = []  # 输出节点列表for node in model.graph.node:if node.name == node_name:start_removal = Trueif start_removal:nodes_to_remove.add(node.name)else:nodes_to_keep.append(node)output.extend(node.output)for node in model.graph.value_info:if node.name in output:shape = [dim.dim_value if (dim.dim_value > 0 and dim.HasField('dim_value')) else Nonefor dim in node.type.tensor_type.shape.dim]output_tensor = helper.make_tensor_value_info(node.name,onnx.TensorProto.FLOAT,shape)model.graph.output.append(output_tensor)new_graph = helper.make_graph(nodes_to_keep,model.graph.name,model.graph.input,[output for output in model.graph.output if output.name not in nodes_to_remove],model.graph.initializer,)new_model = helper.make_model(new_graph, producer_name=model.producer_name)new_model = shape_inference.infer_shapes(new_model)return new_modeldef preprocess_image(image_path, target_shape):"""加载并预处理图像。参数:image_path (str): 图像文件路径。target_shape (tuple): 目标形状 (宽, 高)。返回:np.ndarray: 预处理后的图像数组。"""im = cv2.imread(image_path)im = cv2.resize(im, target_shape)im = im.transpose((2, 0, 1))[::-1]  # HWC 转 CHW, BGR 转 RGBreturn np.ascontiguousarray(im)def convert_onnx_to_rknn(onnx_model_path, rknn_model_path):"""将 ONNX 模型转换为 RKNN 模型。参数:onnx_model_path (str): ONNX 模型路径。rknn_model_path (str): 转换后的 RKNN 模型路径。"""rknn = RKNN()# 加载 ONNX 模型print('--> Loading model')ret = rknn.load_onnx(model=onnx_model_path)if ret != 0:print('Load ONNX model failed!')returnprint('done')# 配置模型print('--> Building model')ret = rknn.build(do_quantization=False)if ret != 0:print('Build RKNN model failed!')returnprint('done')# 导出 RKNN 模型print('--> Export RKNN model')ret = rknn.export_rknn(rknn_model_path)if ret != 0:print('Export RKNN model failed!')returnprint('done')def load_and_run_rknn_model(rknn_model_path, input_data):"""加载 RKNN 模型并运行推理。参数:rknn_model_path (str): RKNN 模型路径。input_data (np.ndarray): 输入数据。返回:list: RKNN 模型的输出结果。"""rknn = RKNN()# 加载 RKNN 模型print('--> Loading RKNN model')ret = rknn.load_rknn(rknn_model_path)if ret != 0:print('Load RKNN model failed!')return []print('done')# 初始化 RKNN 模型print('--> Init runtime environment')ret = rknn.init_runtime()if ret != 0:print('Init runtime environment failed!')return []print('done')# 运行推理print('--> Running model')outputs = rknn.inference(inputs=[input_data])print('done')rknn.release()return outputsdef compare_outputs(onnx_outputs, rknn_outputs, threshold=0.8):"""比较 ONNX 和 RKNN 模型的输出结果。参数:onnx_outputs (list): ONNX 模型的输出结果。rknn_outputs (list): RKNN 模型的输出结果。threshold (float): 准确率阈值。返回:bool: 如果准确率低于阈值，则返回 False，否则返回 True。"""# 计算准确率 (这里假设是简单的相对误差)accuracy = np.mean([np.allclose(onnx_out, rknn_out, rtol=threshold) for onnx_out, rknn_out in zip(onnx_outputs, rknn_outputs)])return accuracy >= thresholddef main():model_path = 'yolov5s.onnx'model = onnx.load(model_path)dtype_map = {'tensor(float)': np.float32,'tensor(int32)': np.int32,'tensor(int64)': np.int64,}all_node_names = get_all_node_names(model)im = preprocess_image('./sample.jpg', (640, 640))for node_name in all_node_names:if node_name:new_model = remove_node_and_following(model, node_name)new_model_path = 'test.onnx'onnx.save(new_model, new_model_path)sess = rt.InferenceSession(new_model_path)output_names = [output.name for output in sess.get_outputs()]input_data = {input.name: im.reshape(input.shape).astype(dtype_map[input.type])for input in sess.get_inputs()}onnx_outputs = sess.run(output_names, input_data)# 转换 ONNX 模型为 RKNN 模型rknn_model_path = 'test.rknn'convert_onnx_to_rknn(new_model_path, rknn_model_path)# 加载并运行 RKNN 模型rknn_outputs = load_and_run_rknn_model(rknn_model_path, im)# 比较 ONNX 和 RKNN 模型的输出结果if not compare_outputs(onnx_outputs, rknn_outputs):print(f'Node {node_name} is the incorrect operator in RKNN model.')breakif __name__ == "__main__":main()