自定义数据训练的rknn模型部署踩坑记录

news2024/11/25 4:32:30

自己训练了一个只有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()