下面是添加了详细注释的优化代码:
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。你可以通过以下命令安装:
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 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()