自己训练了一个只有2种类别的yolov8模型之后，部署到瑞芯微RKNN。
踩坑一：类别的变化
之前用COCO数据集训练的.pt模型转rknn后，output0的shape为8400 x 176.
而把自定义数据集训练的模型转rknn后，output0的shape为8400 x 98.
为什么会不一样？一度以为模型或代码有问题，拿图片测了下发现模型没问题，能检测出来。
分析一下8400x176，其中8400是proposal的个数，176是box相关的64+80个类别+32的mask coefficient.
现在类别变成了2，自然要变成64+2+32=98.
所以，如果后处理中用到类别数=80，画图中的类别标签用了80个的，这里需要修改一下。
踩坑二：模型转化（.pt 转 .rknn）
用转化后的rknn模型检测，只能检测出一种类别，另一种检测不出来。
于是从后处理出发找原因，
现输入图片中有两种物体，每种一个，理论上应该检测出2个物体。
先看NMS之前检测了多少个物体。
std::vector<int> picked;
nms_sorted_bboxes(proposals, picked); //picked里面保存的是proposals的下标
int count = picked.size();
这里NMS之前 picked.size() = 1, 说明NMS之前就只检测出一个，已经不对了。
回到NMS之前的处理，
这里选择较大概率的类别作为label，保存label, score和box坐标。
int label = -1;
float score = -FLT_MAX;
//找到最大score和对应的label
for (int k = 0; k < num_class; k++) {
float confidence = deqnt_affine_to_f32(score_ptr[k], zp,
scale); //反量化,int转float
if (confidence > score) {
label = k;
score = confidence;
}
}
float box_prob = sigmoid(score);
if (box_prob >= thres) {
//忽略一段处理过程
obj.rect.x = x0;
obj.rect.y = y0;
obj.rect.width = x1 - x0;
obj.rect.height = y1 - y0;
obj.label = label;
obj.prob = box_prob;
}
于是追踪label = 0和label = 1的检测概率。
用同一图片测试，很神奇的是.pt模型（转rknn之前）测试中，label=0和label=1的目标 检测概率都>0.9，
而转rknn之后label=1的目标 检测概率只有0.5，label=0的目标 检测概率小得出奇。
观察到score作softmax运算前，label=0的目标 反量化的检测概率小于-7.
认为这是反量化出现了问题，不过反量化用的是从rknn模型中读出的量化参数scales和zp。
for (int i = 0; i < io_num.n_output; ++i) {
output_attrs[i].index = i;
if(rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr)) < 0) {
LOGE("rknn_query output_attrs[%d] fail!\n", i);
return;
}
// set out_scales/out_zps for post_process
out_scales.push_back(output_attrs[i].scale);
out_zps.push_back(output_attrs[i].zp);
}
如果量化参数有问题，那就是转rknn模型的过程有问题。
找到转rknn模型的代码。
OUT_DIR = "rknn_models"
RKNN_MODEL_PATH = '{}/{}.rknn'.format(OUT_DIR,exp)
if NEED_BUILD_MODEL:
DATASET = './dataset.txt'
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform="rkXXXX")
# Load model
print('--> Loading model')
ret = rknn.load_onnx(MODEL_PATH)
if ret != 0:
print('load model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset=DATASET)
if ret != 0:
print('build model failed.')
exit(ret)
print('done')
# Export rknn model
if not os.path.exists(OUT_DIR):
os.mkdir(OUT_DIR)
print('--> Export RKNN model: {}'.format(RKNN_MODEL_PATH))
ret = rknn.export_rknn(RKNN_MODEL_PATH)
if ret != 0:
print('Export rknn model failed.')
exit(ret)
print('done')
else:
ret = rknn.load_rknn(RKNN_MODEL_PATH)
rknn.release()
注意到rknn模型的转换用到了dataset.txt，它指定的是如下图片。
这是原版COCO训练下的yolov8的图片，而现在的自定义数据集中没有这样的图片。
换成现在自定义数据集中的图片，再次转rknn模型。
然后观察scales和zp这两个量化值，发现跟之前不一样了。
再次检测，成功检测出2个目标，且检测概率>0.9，和.pt模型差不多。
总结：
转rknn模型时用的图片可能需要用现有数据集中的图片，不要用不相关的图片。
（原理不清楚）