|
|
最近我在鲁班猫4上部署自己的模型时遇到了rknn转换问题 用的是ppliteseg模型 具体是在开发板上用转换好的rknn执行二分类问题时将整个图片判定为前景 但在pc上使用onnx推理时是正常的（已验证） 转换时用的是官方代码改的 预处理已经检查过没有问题 开发板驱动版本也没问题 用的toolkit2版本都是2.3.2 后面我考虑可能是量化数据集的问题 请问我还需要排查什么问题 转换的代码在后面 刚刚试了下取消量化 发现整个图片都被判为背景了 和之前的结果相反
import sys
import os
import logging
from rknn.api import RKNN
# 配置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# 默认配置
DATASET_PATH = r'/home/lixuboy/data/dataset.txt'
DEFAULT_RKNN_PATH = 'pp_liteseg111.rknn'
DEFAULT_PLATFORM = 'rk3588' # 默认平台
DEFAULT_QUANT = True
SUPPORTED_PLATFORMS = ['rk3562', 'rk3566', 'rk3568', 'rk3576', 'rk3588', 'rv1126b', 'rv1109', 'rv1126', 'rk1808']
SUPPORTED_DTYPES = {
('rk3562', 'rk3566', 'rk3568', 'rk3576', 'rk3588', 'rv1126b'): ['i8', 'fp'],
('rv1109', 'rv1126', 'rk1808'): ['u8', 'fp']
}
def parse_arguments():
"""解析命令行参数并验证"""
if len(sys.argv) < 2:
logger.error("使用方法: python3 {} onnx_model_path [platform] [dtype] [output_rknn_path]".format(sys.argv[0]))
logger.error("平台支持: {}".format(', '.join(SUPPORTED_PLATFORMS)))
logger.error("数据类型支持:")
for platforms, dtypes in SUPPORTED_DTYPES.items():
logger.error(" {}: {}".format(', '.join(platforms), ', '.join(dtypes)))
sys.exit(1)
# 必需参数
model_path = sys.argv[1]
if not os.path.exists(model_path):
logger.error(f"模型文件不存在: {model_path}")
sys.exit(1)
# 可选参数 - 平台
platform = DEFAULT_PLATFORM
if len(sys.argv) > 2:
platform = sys.argv[2]
if platform not in SUPPORTED_PLATFORMS:
logger.error(f"不支持的平台: {platform}")
sys.exit(1)
# 可选参数 - 数据类型
do_quant = DEFAULT_QUANT
if len(sys.argv) > 3:
dtype = sys.argv[3]
# 确定该平台支持的数据类型
valid_dtypes = next((dtypes for platforms, dtypes in SUPPORTED_DTYPES.items() if platform in platforms), [])
if not valid_dtypes:
logger.error(f"未知平台: {platform}")
sys.exit(1)
if dtype not in valid_dtypes:
logger.error(f"平台 {platform} 不支持数据类型: {dtype}")
logger.error(f"支持的数据类型: {', '.join(valid_dtypes)}")
sys.exit(1)
do_quant = dtype in ['i8', 'u8']
# 可选参数 - 输出路径
output_path = DEFAULT_RKNN_PATH
if len(sys.argv) > 4:
output_path = sys.argv[4]
# 确保输出目录存在
output_dir = os.path.dirname(output_path)
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
return model_path, platform, do_quant, output_path
def convert_onnx_to_rknn(model_path, platform, do_quant, output_path):
"""将ONNX模型转换为RKNN模型"""
try:
# 创建RKNN对象
rknn = RKNN(verbose=False)
# 配置模型
logger.info('--> 配置模型')
rknn.config(
mean_values=[[0.485*255, 0.456*255, 0.406*255]],
std_values=[[0.229*255, 0.225*255, 0.225*255]],
target_platform=platform
)
logger.info('模型配置完成')
# 加载模型
logger.info(f'--> 加载ONNX模型: {model_path}')
ret = rknn.load_onnx(model=model_path)
if ret != 0:
logger.error('加载模型失败!')
return False
logger.info('模型加载完成')
# 构建模型
logger.info('--> 构建RKNN模型')
ret = rknn.build(do_quantization=do_quant, dataset=DATASET_PATH)
if ret != 0:
logger.error('构建模型失败!')
return False
logger.info('模型构建完成')
# 导出RKNN模型
logger.info(f'--> 导出RKNN模型到: {output_path}')
ret = rknn.export_rknn(output_path)
if ret != 0:
logger.error('导出RKNN模型失败!')
return False
logger.info('RKNN模型导出完成')
# 释放资源
rknn.release()
return True
except Exception as e:
logger.error(f'转换过程中发生错误: {str(e)}')
return False
if __name__ == '__main__':
# 解析参数
model_path, platform, do_quant, output_path = parse_arguments()
# 打印转换配置
logger.info('===== 转换配置 =====')
logger.info(f'ONNX模型: {model_path}')
logger.info(f'目标平台: {platform}')
logger.info(f'是否量化: {do_quant}')
logger.info(f'输出路径: {output_path}')
logger.info('====================')
# 执行转换
if convert_onnx_to_rknn(model_path, platform, do_quant, output_path):
logger.info('模型转换成功!')
else:
logger.error('模型转换失败!')
sys.exit(1)
|
|
收藏
闽公网安备 35010202001061号