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

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

import cv2
import numpy as np
import onnx
import onnxruntime as rt
from onnx import helper, shape_inference

def 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 = True
        if 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 None
                for 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_model

def 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 转 RGB
    return 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。你可以通过以下命令安装：

2. 补充代码

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

import cv2
import numpy as np
import onnx
import onnxruntime as rt
from onnx import helper, shape_inference
from rknn.api import RKNN

def 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 = True
        if 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 None
                for 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_model

def 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 转 RGB
    return 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!')
        return
    print('done')

    # 配置模型
    print('--> Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build RKNN model failed!')
        return
    print('done')

    # 导出 RKNN 模型
    print('--> Export RKNN model')
    ret = rknn.export_rknn(rknn_model_path)
    if ret != 0:
        print('Export RKNN model failed!')
        return
    print('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 outputs

def 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 >= threshold

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()
            }

            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.')
                break

if __name__ == "__main__":
    main()