在之前的文章中曾讲解过ONNX的量化过程，这里在量化为onnx-fp32格式后，需要部署在开发板RK3588上时，为了更好的适配RK3588的框架格式，我们通过还需进一步量化为RKNN格式，这篇文章便是来讲解RKNN静态量化的过程。
检查ONNX格式
在正式开始RKNN静态量化之前需要对ONNX格式进行检查。
检查路径
onnx_path = Path(onnx_path) if not onnx_path.is_file(): raise FileNotFoundError(f"ONNX模型不存在: {onnx_path}")
检查ONNX模型
这一步是为了检查ONNX本身是否损坏或者导出不完整。
model = onnx.load(str(onnx_path)) onnx.checker.check_model(model)通过onnx.load导入ONNX模型，然后使用onnx.checker.check_model(model)来检查ONNX模型中节点连接、张量定义、算子结构是否符合ONNX规范。
检查输入
这里需要注意的是model.graph.input和模型的真实输入之间的区别，前者实际上是计算图声明的输入，也就是训练时的输入是什么便是什么，但是由于ONNX中定义的问题，因此真实输入的数量可能远小于model.graph.input。这在之前ONNX静态量化中也做过介绍。
initializer_names = {initializer.name for initializer in model.graph.initializer} real_inputs = [ value for value in model.graph.input if value.name not in initializer_names ] # 输入数量 if len(real_inputs) != 1: raise RuntimeError(f"模型输入数量错误，当前输入数量: {len(real_inputs)}") model_input = real_inputs[0] input_shape = get_tensor_shape(model_input) expected_shape = (1, 3, imgsz[0], imgsz[1]) # 输入尺寸 if input_shape != expected_shape: raise ValueError(f"ONNX输入尺寸为{input_shape}，但实际输入应为{expected_shape}")
检查ONNX输入数据类型是否为FP32
通常来说转化为嵌入式端对应的格式都是从ONNX-fp32转化的，虽然int8也可以转化，但是int8对本身精度便有所下降，同时如果后续你想要RKNN-fp16格式，通过ONNX-int8便无法转化。
而静态量化后ONNX的fp32和int8之间最主要的区别除了模型结构、参数之间类型的差别外，还有便是ONNX的int8格式相较于fp32多了伪量化节点(QDQ)，因此这里便有两种检查方式。
# 通过模型结构差别检查 model_input = real_inputs[0] if model_input.type.tensor_type.elem_type != TensorProto.FLOAT: raise ValueError("原始模型为INT8数据类型") # 通过有无QDQ检查 quantized_nodes = [ node for node in model.graph.node if node.op_type in {"QuantizeLinear", "DequantizeLinear"} ] if quantized_nodes: raise ValueError("原始模型为INT8数据类型")
检查输出
检查输出主要需要检查其输出维度、输出个数，这里我的输出已经定死了，因此如果后续变换模型时，这一步便需要微调。
# 输出检查 output_value = [] for output in model.graph.output: output_value.append(get_tensor_shape(output)) if len(output_value) == 2 and output_value[0] == (1, 40, 2100) and output_value[1] == (1, 32 ,80 ,80): print(f"ONNX model: {onnx_path.resolve()}") print(f"ONNX input: {model_input.name} {input_shape}") print(f"ONNX outputs: {output_value[0]}、{output_value[1]}") else: raise ValueError(f"原始模型输出有误，目前模型输出为：{output_value}")
导出RKNN
这一步便是与ONNX中导出最大的不同，ONNX的导出流程必须先导出ONNX-FP32格式才能再此基础上导出ONNX-INT8格式或者做一些额外的算子融合导出混合型格式。而导出RKNN的FP16和INT8类型均使用同一个函数，因此可以直接修改函数中的参数使ONNX-FP32直接导出为RKNN-INT8。
其总流程如下：
分配数据集
当目标量化数据类型为INT8时，需要提前分配好数据集。与之前不同的时，这里在对数据进行预处理后，还需要生成一个同于RKNN量化的dataset.txt。
dataset_file.parent.mkdir(parents=True, exist_ok=True) # 创建dataset.txt所在目录 dataset_content = "\n".join(path.as_posix() for path in calibration_paths) # 将校准图像路径依次写入dataset.txt dataset_file.write_text(dataset_content + "\n", encoding="utf-8") # dataset.txt末尾补上换行符
实例化RKNN对象--RKNN
在进行量化之前还需要实例化RKNN对象
from rknn.api import RKNN rknn = RKNN(verbose, verbose_file)verbose:bool型，是否输出RKNN工具链详细的运行日志
verbose_file:str型，日志文件路径，不写则不做保存
设置目标板端配置--rknn.config
rknn.config这个函数也是量化中最需要注意的函数，其决定你量化的流程和结构是什么样的。
ret = rknn.config( mean_values, std_values, target_platform )不进行混合量化则只需要注意前三个即可。
数据预处理
mean_values:list型，设置输入数据预处理中的每个元素需要减去的值
std_values:list型，设置输入数据预处理中的每个元素除以的值
这两个值同时作用于输入的各个元素中
out_values = (in_values - mean_values) / std_values
其值通常为两层的列表，其外层表示第几个输入，内层表示预处理的数值，内层的数值直接按输入的最高维度作用于各个元素。
[ [...], # input0 [...] # input1 ... ]这里假设预处理为图像做归一化
mean_values = [[0, 0, 0]] std_values = [[255, 255, 255]]中间为3个数值，便是因为图像的最高维度即通道数为3--RGB。
设置目标平台
target_platform:部署的模型最后在哪里使用，如"rk3588"、"rk3568"···
由于部署平台中支持的算子、内部优化、量化以及硬件特性会有差异，因此需要手动选择自己所使用的平台。
TARGET_PLATFORM = "rk3588"
加载ONNX模型--rknn.load
rknn.load_onnx(model, inputs, outputs, input_size_list)model:str型，模型地址
input:张量列表，指定哪些张量作为模型输入
output:张量列表，指定哪些张量作为模型输出
input_size_list:列表，指定输入尺寸
量化RKNN模型--rknn.build
rknn.build( do_quantization, dataset, rknn_batch_size )do_quantization:boo型，是否构建INT8类型
dataset:str型，数据集地址
rknn_batch_size:输入批次，通常填入输入的最高维度
如果只需要导出FP16则直接填写do_quantization=false，INT8则除了需要将do_quantization改为true外，还需要填写数据集地址。
导出RKNN--rknn.export_rknn
rknn.export_rknn(export_path)export_path:str型，RKNN模型导出地址
需要注意的是export_rknn本身不进行量化，其功能就是序列化通过build构建的在当前内存中的RKNN模型，然后储存在指定地址。
导出混合RKNN
导出混合RKNN即是大面INT8量化，局部FP16保持，通常进行的混合量化是模型主体INT8量化，输出头FP16保持。
其流程比上面的流程更多，但是仍然维持这上面的主体流程。
分配数据集
同上
设置目标板端配置
在进行上述的数据预处理和设置目标平台后，还需进行下面的步骤。
查找分割头
在分割头的查找中我们通常需要使用正则表达式，具体介绍可以参考我上一篇文章（ONNX静态量化）。
def find_segment_head_info(onnx_path): """查找分割头""" model = onnx.load(str(onnx_path)) segment_module_indices = set() for node in model.graph.node: if "/proto/" not in node.name: # 仅对掩膜模块进行分类 continue match = re.search( r"(?:^|/)model\.(\d+)(?:/|\.|$)", node.name, ) if match: segment_module_indices.add(int(match.group(1))) if len(segment_module_indices) != 1: raise RuntimeError("无法唯一识别YOLOv8 Segment头模块，找到的候选编号为: "f"{sorted(segment_module_indices)}") module_index = next(iter(segment_module_indices)) # 建立分割头迭代器 module_pattern = re.compile(rf"(?:^|/)model\.{module_index}(?:/|\.|$)") # 查找所有分割头模块 head_nodes = [ node for node in model.graph.node if module_pattern.search(node.name) ] if not head_nodes: raise RuntimeError( f"ONNX中没有找到model.{module_index}的分割头节点" ) output_names = tuple(output.name for output in model.graph.output) print( f"分割头: model.{module_index}, " f"ONNX nodes={len(head_nodes)}, outputs={output_names}" ) return module_index, output_names
设置配置
ret = rknn.config( mean_values, std_values, target_platform, quantized_dtype, quantized_algorithm, quantized_method, float_dtype, optimization_level, op_target, single_core_mode, enable_flash_attention, compress_weight )除了上面的mean_values, std_values, target_platform这三个参数外，还需要填写下述参数。
quantized_dtype:权重和激活量化位宽
quantized_dtype 含义 支持平台 w8a8权重 INT8 非对称量化，激活 INT8 非对称量化 除 RK2118 外均支持 w4a16权重 4 bit，激活 FP16 仅 RK3576、RV1126B w8a16权重 INT8，激活 FP16 仅 RK3562 w16a16i权重和激活均为 16 bit 非对称整数量化 仅 RV1103、RV1106 w16a16i_dfp权重和激活均为 16 bit 动态定点 仅 RV1103、RV1106 quantized_algorithm:量化校准算法
|---------------------|---------------|---------|-------------|
| quantized_algorithm | 含义 | 建议样本数 | 适用场景 |
| normal | 量化标准算法 | 20~100 | 首次转化、较为常见模型 |
| mmse | 按最小均方误差搜索量化范围 | 20~50 | 对数值误差敏感的模型 |
| kl_divergense | 按KL散度选择范围 | 20~100 | 激活分布不均匀、长尾等 |quantized_method:权重量化粒度
|------------------|------------------------------|----------------|
| quantized_method | 含义 | 适用场景 |
| layer | 整个权重张量全部使用scale+zero_point量化 | 权重分布均匀，对模型大小敏感 |
| channel | 每个输出通道分别计算量化参数 | 精度更高 |float_dtype:未被INT8量化的层保持上面数据类型，通常为"float16"
optimization_level:图优化强度
|--------------------|---------------|
| optimization_level | 含义 |
| 3 | 开启全部默认优化 |
| 2 | 关闭部分可能影响精度的优化 |
| 1 | 关闭更多优化 |
| 0 | 基本关闭图优化 |op_target:手动指定某些算子运行在CPU/NPU/GPU上
# 按输出张量名指定 op_target={ "/encoder/block0/Add_output_0": "cpu", } # 按算子类型指定 op_target={ "Add": "cpu", "Softmax": "npu", }single_core_mode:bool型，是否只生成RK3588单NPU核模型，通常在同时运行多个模型时，且该模型较小、简单时为true
enable_flash_attention:bool型，是否启动Flash Attention优化
在模型结构由注意力机制时，令其为true，可以减少Attention对内存带宽的访问，同时利用换成进行分块计算，并有可能加速长序列Attention，但是可能使得RKNN增大。
compress_weight:bool型，是否压缩模型权重
当为false时为最常规的转化路径，适合作为对比基线，但是当还有大量Conv/Linear/MatMul权重的模型、大型Transformor等本身模型文件过大，从而导致储存空间受限的模型，需要适用true。
因此这里我们需要重新设置配置文件
ret = rknn.config( mean_values=MEAN_VALUES, std_values=STD_VALUES, target_platform=target_platform, quantized_dtype="w8a8", quantized_algorithm="normal", quantized_method="channel", float_dtype=SEGMENT_HEAD_FLOAT_DTYPE, )
加载ONNX模型
同上
第一步混合量化--rknn.hybrid_quantization_step1
其本质的作用是先完成一次默认INT8校准，再生成可供人工指定高精度层的中间模型和配置。
rknn.hybrid_quantization_step1( dataset, rknn_batch_size, proposal, proposal_dataset_size, custom_hybrid )dataset:str型，校准集地址
rknn_batch_size:模型批次
proposal:bool型，是否分析并推荐需要保留高精度的层，当为true时，其可能回推荐过多层数，从而导致FP16过多，但是当为false时，其所有层的类型均为INT8，后续便需要自己编辑。
proposal_dataset_size:推荐分析使用多少条校准数据，为上一参数服务，其越大精度越高，但是过大会导致分析时间明显变长同时结果也未必会改善。
custom_hybrid:指定需要进行混合量化分析的子图
ret = rknn.hybrid_quantization_step1( dataset=str(dataset_path), rknn_batch_size=1, proposal=False, )
判断模型缓存文件是否正确生成
生成的文件通常和onnx模型名字相同，仅更换了后缀名
# 判断临时缓存模型文件是否生成 stem = onnx_path.stem model_input = work_dir / f"{stem}.model" data_input = work_dir / f"{stem}.data" raw_cfg = work_dir / f"{stem}.quantization.cfg" for artifact in (model_input, data_input, raw_cfg): if not artifact.is_file() or artifact.stat().st_size == 0: raise RuntimeError( f"混合量化第一步没有生成有效文件: {artifact}" )
模型缓存文件
临时模型文件--xxx.model
临时模型文件主要保持经过RKNN-Toolkit2解析和优化的中间模型结构，如优化后的计算图、算子和张量之间的连接、张量名称以及第二步混合量化重建模型需要的编译信息等。
由于第一步混合量化已经进行了常量折叠、算子融合等步骤，因此此时生成的临时模型文件和onnx保存的模型文件存在差异。
二进制数据文件--xxx.data
临时模型文件配套的二进制数据文件，其通常包括模型常量和权重数据、第一步混合量化的模型数据等
配置文件--xxx.quantization.cfg
这是最需要了解的文件，也是后续主要需要修改的文件，其主要作用便是提供给用户查看和修改模型格式、配置等的文件。
其通常包含下述两个信息：
对应层精度参数--custom_quantize_layers
该参数用于手动指定某个输出张量对应的层使用的精度类型，如
custom_quantize_layers: /model.4/Add_output_0: float16 /model.8/Softmax_output_0: float16表示除了产生/model.4/Add_output_0和/model.8/Softmax_output_0的两层保持FP16，其余层均保持INT8。
而/model.4/Add_output_0和/model.8/Softmax_output_0即是模型中的两个张量，而保持FP16的便是生成这两个张量的计算和输出。输入仍然是INT8，但是在计算时通过反量化为FP16，从而进行的FP16计算，同时输出FP16，但是在其下一层仍然需要量化为INT8计算。
对应层量化参数--quantize_parameters
其通常内容如下：
quantize_parameters: /model.0/Conv_output_0: qtype: asymmetric_quantized qmethod: layer dtype: int8 min: - -12.5 max: - 18.7 scale: - 0.122 zero_point: - -26其表示在产生/model.0/Conv_output_0这一层的参数信息，这个参数不易改动，其和上面两个文件之间相对应，修改该参数容易造成该配置文件无法和上面的两个文件匹配，从而导致最后第二步的混合量化出错。
加载临时ONNX计算图并修改cfg配置
加载临时ONNX计算图
在上面我们已经讲述了xxx.model这个文件是第一步混合量化产生的临时模型文件。这个模型文件仍然是通过ONNX来进行加载的。
optimized_model = onnx.load( str(optimized_model_path), load_external_data=False, )
排除常量算子
由于后续查找分割头中会出现常量算子，因此在这里我们可以通过算子类型得到排除常量算子的所有算子。
# 排除常量算子 tensor_producers = { tensor_name: node_index for node_index, node in enumerate(optimized_model.graph.node) if node.op_type != "Constant" for tensor_name in node.output }
读取cfg配置文件
# 读取cfg配置文件 lines = raw_cfg_path.read_text( encoding="utf-8-sig" ).splitlines() try: # 精度参数起始地址 custom_index = next( index for index, line in enumerate(lines) if line.strip().startswith("custom_quantize_layers:") ) # 量化参数起始地址 parameters_index = next( index for index, line in enumerate(lines) if line.strip() == "quantize_parameters:" ) except StopIteration as error: raise RuntimeError( "RKNN混合量化配置缺少custom_quantize_layers或quantize_parameters段" ) from error # 精度参数在量化参数上方 if custom_index >= parameters_index: raise RuntimeError("RKNN混合量化配置段落顺序异常")
获取临时文件分割头
# 获取临时文件分割头 tensor_names = [] selected_producers = set() for line in lines[parameters_index + 1:]: if line and not line[0].isspace(): break match = re.match(r"^ {4}(.+):\s*$", line) if not match: continue tensor_name = decode_yaml_key(match.group(1)) # 将所有算子全部去掉引号 # 检查算子是否和之前得到的分割头层级相同 if is_segment_head_tensor( tensor_name, module_index=module_index, output_names=output_names, ): producer_index = tensor_producers.get(tensor_name) # 配置中还可能包含已被常量折叠移除的张量，以及后续自动插入的int8输出转换张量 if producer_index is None: continue # Split/Slice等多输出算子只需指定一次 if producer_index in selected_producers: continue selected_producers.add(producer_index) tensor_names.append(tensor_name)这是这里这里最为主要的一步，首先遍历所有量化参数得到模型中每一个算子的名称，其格式固定为：4个空格/model.x/xxx，如下
quantize_parameters: images: qtype: asym qmethod: layer dtype: float32 min: - 0.0 max: - 1.0 scale: [] zero_point: [] ori_min: - 0.0 ori_max: - 1.0 /model.0/conv/Conv_output_0: qtype: asym qmethod: layer dtype: int8 min: - -63.26831817626953 max: - 65.14521789550781 scale: - 0.5035824943991268 zero_point: - -2 ori_min: - -63.26831817626953 ori_max: - 65.14521789550781 ....因此这里可以再次使用正则表达式进行查找，得出的结构全部使用去掉引号进行保存
def decode_yaml_key(token): """解析RKNN量化配置中的简单YAML键""" token = token.strip() # 处理双引号 if len(token) >= 2 and token[0] == token[-1] == '"': return json.loads(token) # 处理单引号 if len(token) >= 2 and token[0] == token[-1] == "'": return token[1:-1].replace("''", "'") return token由于在前面查找分割头中得到了分割头所在的层级，因此可以通过判断是否为同一层级的算子来得出是否为分割头
def is_segment_head_tensor(tensor_name, module_index, output_names): """判断RKNN优化后的张量是否属于Segment头或其最终输出""" module_path_marker = f"/model.{module_index}/" module_weight_marker = f"model.{module_index}." if module_path_marker in tensor_name or module_weight_marker in tensor_name: return True for output_name in output_names: if tensor_name == output_name: return True if tensor_name.startswith( ( f"{output_name}-", f"{output_name}_", f"{output_name}#", ) ): return True return False虽然得出了改算子为分割头的一部分，但是并没有区分是否为常值算子，因此可以通过之前得到的排除了常值算子后的字典进行判断，在字典中的便不是常值算子，否则便是。同时再通过集合不可重复的特性排除重复的算子，进而得出最终需要保持精度的算子名称。
修改cfg配置
其保存格式应该为：4个空格+算子名称+":"+类型
# 修改配置文件 custom_lines = ["custom_quantize_layers:"] custom_lines.extend( " " # 4个空格 + json.dumps(tensor_name, ensure_ascii=False) # 算子名称，需要加上双引号 + f": {float_dtype}" # 冒号+类型 for tensor_name in tensor_names ) modified_lines = ( lines[:custom_index] # 保持类型参数前面的信息 + custom_lines # 添加的类型参数信息 + lines[parameters_index:] # 保持精度参数后面的信息 ) output_cfg_path.parent.mkdir(parents=True, exist_ok=True) output_cfg_path.write_text( "\n".join(modified_lines) + "\n", encoding="utf-8", )
其实这里实现的办法有很多。
例如你可以通过再得到cfg配置文件后，通过正则表达式查找到所有需要修改的算子名称，然后直接输出，通过手动将这些算子全部更改配置。
当然也可以使用前面我们提到的第一步量化中的custom_hybrid参数，该参数比直接修改更加方便，因为其可以只用边界进行修改，也就是说可以将分割头最初的输入和最终的输出通过列表的形式赋值，无需再依次写这出所有的算子名称。
第二步混合量化--rknn.hybrid_quantization_step2
rknn.hybrid_quantization_step2( model_input="model_name.model", data_input="model_name.data", model_quantization_cfg="model_name.quantization.cfg", )
参数 含义 model_input第一步混合量化生成的临时模型 data_input第一步混合量化生成的校准、量化相关数据 model_quantization_cfg修改后的逐层量化配置
判断模型缓存文件是否正确生成
# 判断临时缓存模型文件是否生成 stem = onnx_path.stem model_input = work_dir / f"{stem}.model" data_input = work_dir / f"{stem}.data" raw_cfg = work_dir / f"{stem}.quantization.cfg" for artifact in (model_input, data_input, raw_cfg): if not artifact.is_file() or artifact.stat().st_size == 0: raise RuntimeError( f"混合量化第一步没有生成有效文件: {artifact}" )
导出RKNN模型
同上
总结
这篇文章主要讲解的是量化数据类型的方法，如果需要进行手动的常值折叠、算子融合等均需要在检查ONNX格式之前进行，完成后再进行进一步的量化。
完整代码
import json import os import re from pathlib import Path import cv2 import onnx from onnx import TensorProto from rknn.api import RKNN # 项目/模型/图像地址 PROJECT_ROOT = Path(__file__).resolve().parent ONNX_PATH = Path(r"E:\code\python\machine\runs\segment\third_tune\weights\best-seg-third-fp32.onnx") CALIB_DIR = Path(r"E:\code\python\machine\data\third_tune\data\calibration") # 保存路径 RKNN_OUTPUT_DIR = PROJECT_ROOT / "runs" / "segment" / "third_tune" / "weights" FP32_RKNN_OUTPUT = RKNN_OUTPUT_DIR / "best-seg-third-fp32.rknn" INT8_RKNN_OUTPUT = RKNN_OUTPUT_DIR / "best-seg-third-int8.rknn" CALIB_CACHE_DIR = PROJECT_ROOT / "calibration_output" / "third_tune_rknn_320x320" CALIB_DATASET_FILE = CALIB_CACHE_DIR / "dataset.txt" HYBRID_WORK_DIR = CALIB_CACHE_DIR / "hybrid_int8_fp16" HYBRID_HEAD_FP16_CFG = ( HYBRID_WORK_DIR / f"{ONNX_PATH.stem}.segment_head_fp16.quantization.cfg" ) # 初始化参数 TARGET_PLATFORM = "rk3588" # 开发板实际芯片型号 HEIGHT = 320 # 图像高 WIDTH = 320 # 图像宽 MAX_CALIB_IMAGES = None # None表示使用全部校准集图像 VERBOSE = True # 显示RKNN Toolkit详细转换日志 SEGMENT_HEAD_FLOAT_DTYPE = "float16" # RK3588浮点NPU层使用FP16 MIN_SEGMENT_HEAD_TENSORS = 20 # 防止层名识别失败后误生成全INT8模型 # 原ONNX输入为RGB float32且数值范围为0~1。 # 把除以255的操作写入RKNN模型后，部署端应输入RGB uint8图像，不要再次除以255。 MEAN_VALUES = [[0.0, 0.0, 0.0]] STD_VALUES = [[255.0, 255.0, 255.0]] IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} def normalize_imgsz(height, width): """检查图像缩放尺寸是否正确""" if isinstance(height, int) and isinstance(width, int): return height, width raise ValueError("图像高度和宽度必须是正整数") def list_images(image_dir, max_images): """获取所有校准图像地址""" image_dir = Path(image_dir) if not image_dir.is_dir(): raise FileNotFoundError(f"校准图像文件夹不存在: {image_dir}") if max_images is not None and ( not isinstance(max_images, int) or max_images <= 0 ): raise ValueError("MAX_CALIB_IMAGES必须是正整数或None。") paths = sorted( path for path in image_dir.rglob("*") if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES ) if not paths: raise FileNotFoundError(f"没有找到校准图像: {image_dir}") return paths[:max_images] if max_images else paths def letterbox(image, new_shape): """保resize + padding""" src_h, src_w = image.shape[:2] dst_h, dst_w = new_shape scale = min(dst_h / src_h, dst_w / src_w) resized_w = int(round(src_w * scale)) resized_h = int(round(src_h * scale)) pad_w = dst_w - resized_w pad_h = dst_h - resized_h if (src_w, src_h) != (resized_w, resized_h): image = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR,) left = int(round(pad_w / 2 - 0.1)) right = int(round(pad_w / 2 + 0.1)) top = int(round(pad_h / 2 - 0.1)) bottom = int(round(pad_h / 2 + 0.1)) return cv2.copyMakeBorder( image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114), ) def prepare_calibration_dataset(image_paths, cache_dir, dataset_file, imgsz): """生成letterbox校准图像和RKNN要求的dataset.txt""" cache_dir = Path(cache_dir) dataset_file = Path(dataset_file) image_cache_dir = cache_dir / "images" image_cache_dir.mkdir(parents=True, exist_ok=True) calibration_paths = [] image_count = len(image_paths) for index, image_path in enumerate(image_paths, start=1): image = cv2.imread(str(image_path)) if image is None: raise ValueError(f"无法读取校准图像: {image_path}") image = letterbox(image, imgsz) output_path = image_cache_dir / f"{index:06d}.jpg" saved = cv2.imwrite( str(output_path), image, [cv2.IMWRITE_JPEG_QUALITY, 95], ) if not saved: raise OSError(f"无法保存校准图像: {output_path}") calibration_paths.append(output_path.resolve()) if index % 50 == 0 or index == image_count: print(f"正在准备校准集: {index}/{image_count}") dataset_file.parent.mkdir(parents=True, exist_ok=True) # 创建dataset.txt所在目录 dataset_content = "\n".join(path.as_posix() for path in calibration_paths) # 将校准图像路径依次写入dataset.txt dataset_file.write_text(dataset_content + "\n", encoding="utf-8") # dataset.txt末尾补上换行符 return dataset_file def get_tensor_shape(value_info): """读取ONNX张量的固定维度""" dimensions = [] for dimension in value_info.type.tensor_type.shape.dim: if dimension.HasField("dim_value"): dimensions.append(dimension.dim_value) elif dimension.HasField("dim_param"): dimensions.append(dimension.dim_param) else: dimensions.append(None) return tuple(dimensions) def validate_onnx_model(onnx_path, imgsz): """检查ONNX模型、输入精度和静态输入尺寸""" # 路径检查 onnx_path = Path(onnx_path) if not onnx_path.is_file(): raise FileNotFoundError(f"ONNX模型不存在: {onnx_path}") # 模型检查 model = onnx.load(str(onnx_path)) onnx.checker.check_model(model) # 输入检查 # 输入数量检查 initializer_names = {initializer.name for initializer in model.graph.initializer} real_inputs = [ value for value in model.graph.input if value.name not in initializer_names ] if len(real_inputs) != 1: raise RuntimeError(f"模型输入数量错误，当前输入数量: {len(real_inputs)}") # 输入尺寸检查 model_input = real_inputs[0] input_shape = get_tensor_shape(model_input) expected_shape = (1, 3, imgsz[0], imgsz[1]) if input_shape != expected_shape: raise ValueError(f"ONNX输入尺寸为{input_shape}，但实际输入应为{expected_shape}") # 原始模型检查 # 模型结构检查 if model_input.type.tensor_type.elem_type != TensorProto.FLOAT: raise ValueError("原始模型为INT8数据类型") # QDQ节点检查 quantized_nodes = [ node for node in model.graph.node if node.op_type in {"QuantizeLinear", "DequantizeLinear"} ] if quantized_nodes: raise ValueError("原始模型为INT8数据类型") # 输出检查 output_value = [] for output in model.graph.output: output_value.append(get_tensor_shape(output)) if len(output_value) == 2 and output_value[0] == (1, 40, 2100) and output_value[1] == (1, 32 ,80 ,80): print(f"ONNX model: {onnx_path.resolve()}") print(f"ONNX input: {model_input.name} {input_shape}") print(f"ONNX outputs: {output_value[0]}、{output_value[1]}") else: raise ValueError(f"原始模型输出有误，目前模型输出为：{output_value}") return model_input.name def find_segment_head_info(onnx_path): """查找分割头""" model = onnx.load(str(onnx_path)) segment_module_indices = set() for node in model.graph.node: if "/proto/" not in node.name: # 仅对掩膜模块进行分类 continue match = re.search( r"(?:^|/)model\.(\d+)(?:/|\.|$)", node.name, ) if match: segment_module_indices.add(int(match.group(1))) if len(segment_module_indices) != 1: raise RuntimeError("无法唯一识别YOLOv8 Segment头模块，找到的候选编号为: "f"{sorted(segment_module_indices)}") module_index = next(iter(segment_module_indices)) # 建立分割头迭代器 module_pattern = re.compile(rf"(?:^|/)model\.{module_index}(?:/|\.|$)") # 查找所有分割头模块 head_nodes = [ node for node in model.graph.node if module_pattern.search(node.name) ] if not head_nodes: raise RuntimeError( f"ONNX中没有找到model.{module_index}的分割头节点" ) output_names = tuple(output.name for output in model.graph.output) print( f"分割头: model.{module_index}, " f"ONNX nodes={len(head_nodes)}, outputs={output_names}" ) return module_index, output_names def decode_yaml_key(token): """解析RKNN量化配置中的简单YAML键""" token = token.strip() # 处理双引号 if len(token) >= 2 and token[0] == token[-1] == '"': return json.loads(token) # 处理单引号 if len(token) >= 2 and token[0] == token[-1] == "'": return token[1:-1].replace("''", "'") return token def is_segment_head_tensor(tensor_name, module_index, output_names): """判断RKNN优化后的张量是否属于Segment头或其最终输出""" module_path_marker = f"/model.{module_index}/" module_weight_marker = f"model.{module_index}." if module_path_marker in tensor_name or module_weight_marker in tensor_name: return True for output_name in output_names: if tensor_name == output_name: return True if tensor_name.startswith( ( f"{output_name}-", f"{output_name}_", f"{output_name}#", ) ): return True return False def keep_segment_head_float( raw_cfg_path, output_cfg_path, optimized_model_path, module_index, output_names, float_dtype, ): """修改RKNN混合量化配置，使整个Segment头保持FP16""" raw_cfg_path = Path(raw_cfg_path) output_cfg_path = Path(output_cfg_path) optimized_model_path = Path(optimized_model_path) # 加载临时ONNX计算图 optimized_model = onnx.load( str(optimized_model_path), load_external_data=False, ) # 排除常量算子 tensor_producers = { tensor_name: node_index for node_index, node in enumerate(optimized_model.graph.node) if node.op_type != "Constant" for tensor_name in node.output } # 读取cfg配置文件 lines = raw_cfg_path.read_text( encoding="utf-8-sig" ).splitlines() try: custom_index = next( index for index, line in enumerate(lines) if line.strip().startswith("custom_quantize_layers:") ) parameters_index = next( index for index, line in enumerate(lines) if line.strip() == "quantize_parameters:" ) except StopIteration as error: raise RuntimeError( "RKNN混合量化配置缺少custom_quantize_layers或quantize_parameters段" ) from error if custom_index >= parameters_index: raise RuntimeError("RKNN混合量化配置段落顺序异常") # 获取临时文件分割头 tensor_names = [] selected_producers = set() for line in lines[parameters_index + 1:]: if line and not line[0].isspace(): break match = re.match(r"^ {4}(.+):\s*$", line) if not match: continue tensor_name = decode_yaml_key(match.group(1)) if is_segment_head_tensor( tensor_name, module_index=module_index, output_names=output_names, ): producer_index = tensor_producers.get(tensor_name) # 配置中还可能包含已被常量折叠移除的张量，以及后续自动插入的int8输出转换张量 if producer_index is None: continue # Split/Slice等多输出算子只需指定一次 if producer_index in selected_producers: continue selected_producers.add(producer_index) tensor_names.append(tensor_name) tensor_names = list(dict.fromkeys(tensor_names)) if len(tensor_names) < MIN_SEGMENT_HEAD_TENSORS: raise RuntimeError( f"仅在RKNN配置中识别到{len(tensor_names)}个Segment头张量，少于安全下限{MIN_SEGMENT_HEAD_TENSORS}；已停止，避免误生成全INT8模型" ) missing_outputs = [ output_name for output_name in output_names if not any( tensor_name == output_name or tensor_name.startswith( ( f"{output_name}-", f"{output_name}_", f"{output_name}#", ) ) for tensor_name in tensor_names ) ] if missing_outputs: raise RuntimeError( f"RKNN配置中没有识别到以下Segment模型输出: {missing_outputs}" ) # 修改配置文件 custom_lines = ["custom_quantize_layers:"] custom_lines.extend( " " # 4个空格 + json.dumps(tensor_name, ensure_ascii=False) # 算子名称，需要加上双引号 + f": {float_dtype}" # 冒号+类型 for tensor_name in tensor_names ) modified_lines = ( lines[:custom_index] # 保持类型参数前面的信息 + custom_lines # 添加的类型参数信息 + lines[parameters_index:] # 保持精度参数后面的信息 ) output_cfg_path.parent.mkdir(parents=True, exist_ok=True) output_cfg_path.write_text( "\n".join(modified_lines) + "\n", encoding="utf-8", ) print(f"混合量化配置地址: {output_cfg_path.resolve()}") return output_cfg_path def check_rknn_ret(step, ret): """检查RKNN Toolkit接口返回值""" if ret not in (None, 0): raise RuntimeError(f"{step}失败，RKNN返回码: {ret}") def convert_onnx_to_rknn( rknn, onnx_path, output_path, target_platform, do_quantization, dataset_path, ): """将FP32 ONNX构建为RKNN""" onnx_path = Path(onnx_path) output_path = Path(output_path) if do_quantization and dataset_path is None: raise ValueError("构建INT8 RKNN时必须提供校准集") precision_name = "INT8" if do_quantization else "FP16" print(f"\n转化{precision_name}类型精度") rknn = rknn(verbose = VERBOSE) try: ret = rknn.config( mean_values=MEAN_VALUES, std_values=STD_VALUES, target_platform=target_platform, ) check_rknn_ret("配置模型", ret) ret = rknn.load_onnx(model=str(onnx_path)) check_rknn_ret("加载ONNX模型", ret) build_args = {"do_quantization": do_quantization} if do_quantization: build_args["dataset"] = str(Path(dataset_path).resolve()) ret = rknn.build(**build_args) check_rknn_ret(f"构建{precision_name}模型", ret) output_path.parent.mkdir(parents=True, exist_ok=True) ret = rknn.export_rknn(str(output_path)) check_rknn_ret("导出RKNN模型", ret) finally: rknn.release() if not output_path.is_file() or output_path.stat().st_size == 0: raise RuntimeError(f"RKNN接口未报错，但输出文件无效: {output_path}") print(f"{precision_name} RKNN: {output_path.resolve()}") return output_path def generate_hybrid_quantization_files( rknn, onnx_path, dataset_path, work_dir, target_platform, verbose, ): """执行混合量化第一步""" onnx_path = Path(onnx_path).resolve() dataset_path = Path(dataset_path).resolve() work_dir = Path(work_dir).resolve() if not dataset_path.is_file(): raise FileNotFoundError(f"校准dataset.txt不存在: {dataset_path}") work_dir.mkdir(parents=True, exist_ok=True) rknn = rknn(verbose=verbose) previous_cwd = Path.cwd() try: ret = rknn.config( mean_values=MEAN_VALUES, std_values=STD_VALUES, target_platform=target_platform, quantized_dtype="w8a8", quantized_algorithm="normal", quantized_method="channel", float_dtype=SEGMENT_HEAD_FLOAT_DTYPE, ) check_rknn_ret("配置混合量化模型", ret) ret = rknn.load_onnx(model=str(onnx_path)) check_rknn_ret("加载混合量化ONNX", ret) # Toolkit2会把step1生成的文件写入当前工作目录 os.chdir(work_dir) ret = rknn.hybrid_quantization_step1( dataset=str(dataset_path), rknn_batch_size=1, proposal=False, ) check_rknn_ret("混合量化第一步", ret) finally: os.chdir(previous_cwd) rknn.release() # 判断临时缓存模型文件是否生成 stem = onnx_path.stem model_input = work_dir / f"{stem}.model" data_input = work_dir / f"{stem}.data" raw_cfg = work_dir / f"{stem}.quantization.cfg" for artifact in (model_input, data_input, raw_cfg): if not artifact.is_file() or artifact.stat().st_size == 0: raise RuntimeError( f"混合量化第一步没有生成有效文件: {artifact}" ) return model_input, data_input, raw_cfg def build_hybrid_rknn( rknn, model_input, data_input, quantization_cfg, output_path, verbose, ): """执行混合量化第二步并导出骨干INT8、Segment头FP16的RKNN""" model_input = Path(model_input) data_input = Path(data_input) quantization_cfg = Path(quantization_cfg) output_path = Path(output_path) for artifact in (model_input, data_input, quantization_cfg): if not artifact.is_file(): raise FileNotFoundError(f"混合量化输入不存在: {artifact}") print("\n混合量化RKNN中...") rknn = rknn(verbose=verbose) try: ret = rknn.hybrid_quantization_step2( model_input=str(model_input.resolve()), data_input=str(data_input.resolve()), model_quantization_cfg=str(quantization_cfg.resolve()), ) check_rknn_ret("混合量化第二步", ret) output_path.parent.mkdir(parents=True, exist_ok=True) ret = rknn.export_rknn(str(output_path)) check_rknn_ret("导出混合INT8/FP16 RKNN", ret) finally: rknn.release() if not output_path.is_file() or output_path.stat().st_size == 0: raise RuntimeError( f"RKNN接口未报错，但混合量化输出文件无效: {output_path}" ) print(f"混合量化RKNN文件: {output_path.resolve()}") return output_path def convert_onnx_to_hybrid_rknn( rknn, onnx_path, dataset_path, output_path, work_dir, output_cfg_path, target_platform, ): """仅量化骨干和颈部，Segment检测/Mask/Proto头保持FP16""" module_index, output_names = find_segment_head_info(onnx_path) model_input, data_input, raw_cfg = ( generate_hybrid_quantization_files( rknn=rknn, onnx_path=onnx_path, dataset_path=dataset_path, work_dir=work_dir, target_platform=target_platform, verbose=VERBOSE, ) ) mixed_cfg = keep_segment_head_float( raw_cfg_path=raw_cfg, output_cfg_path=output_cfg_path, optimized_model_path=model_input, module_index=module_index, output_names=output_names, float_dtype=SEGMENT_HEAD_FLOAT_DTYPE, ) return build_hybrid_rknn( rknn=rknn, model_input=model_input, data_input=data_input, quantization_cfg=mixed_cfg, output_path=output_path, verbose=VERBOSE, ) def prepare_validation_data_yaml(source_yaml, dataset_root, output_yaml): """生成Windows和WSL都能使用的验证集YAML""" source_yaml = Path(source_yaml) dataset_root = Path(dataset_root) output_yaml = Path(output_yaml) if not source_yaml.is_file(): raise FileNotFoundError(f"数据集YAML文件未找到: {source_yaml}") if not (dataset_root / "val" / "images").is_dir(): raise FileNotFoundError( f"验证集未找到: {dataset_root / 'val' / 'images'}" ) if not (dataset_root / "val" / "labels").is_dir(): raise FileNotFoundError( f"验证集标签未找到: {dataset_root / 'val' / 'labels'}" ) dataset_path_line = f"path: {dataset_root.resolve().as_posix()}" source_lines = source_yaml.read_text(encoding="utf-8-sig").splitlines() output_lines = [] path_replaced = False for line in source_lines: if not path_replaced and line.startswith("path:"): output_lines.append(dataset_path_line) path_replaced = True else: output_lines.append(line) if not path_replaced: output_lines.insert(0, dataset_path_line) output_yaml.parent.mkdir(parents=True, exist_ok=True) output_yaml.write_text( "\n".join(output_lines) + "\n", encoding="utf-8", ) print(f"数据集YAML文件: {output_yaml.resolve()}") return output_yaml def main(): if not isinstance(TARGET_PLATFORM, str) or not TARGET_PLATFORM.strip(): raise ValueError("请设置正确的TARGET_PLATFORM") imgsz = normalize_imgsz(HEIGHT, WIDTH) validate_onnx_model(ONNX_PATH, imgsz) rknn = RKNN output_paths = [] dataset_path = None # FP16模型导出 output_paths.append( convert_onnx_to_rknn( rknn=rknn, onnx_path=ONNX_PATH, output_path=FP32_RKNN_OUTPUT, target_platform=TARGET_PLATFORM, do_quantization=False, dataset_path=dataset_path ) ) # 数据集处理 image_paths = list_images(CALIB_DIR, MAX_CALIB_IMAGES) dataset_path = prepare_calibration_dataset( image_paths=image_paths, cache_dir=CALIB_CACHE_DIR, dataset_file=CALIB_DATASET_FILE, imgsz=imgsz, ) # 混合INT8/FP16模型导出：骨干和颈部INT8，整个Segment头保持FP16 output_paths.append( convert_onnx_to_hybrid_rknn( rknn=rknn, onnx_path=ONNX_PATH, dataset_path=dataset_path, output_path=INT8_RKNN_OUTPUT, work_dir=HYBRID_WORK_DIR, output_cfg_path=HYBRID_HEAD_FP16_CFG, target_platform=TARGET_PLATFORM ) ) print("\n量化完成:") for output_path in output_paths: print(output_path.resolve()) if __name__ == "__main__": main()