【CV学习笔记】tensorrt篇之HRnet

1、前言

在计算机视觉领域，除了对常见的检测、分割网络进行了学习，最近又看到了关注的up主发布了有关关键点检测的网络，正是本文的所学习的网络:HRnet，本文在源代码的基础上，进行了python版本的后处理优化、并使用cupy和torch分别来进行后处理，并对比了两者在后处理的时间，同时在cpp代码里面实现了cuda后处理加速。

原文地址为:https://blog.csdn.net/qq_37541097/article/details/124346626?spm=1001.2014.3001.5501

个人学习地址为:https://github.com/Rex-LK/tensorrt_learning
欢迎正在学习或者想学的CV的同学进群一起讨论与学习，v:Rex1586662742，q群:468713665

2、HRnet简介

在原hrnet博文里面up主绘制了相应的网络结构图，这里就直接拿来用了。

从图中可看出，图像经过一系列的特征提取，尤其是stage1~3里面的将不同尺度的特征信息进行融合，最终得到了一个64×48×17的特征图，其中64×48表示将图片划分为64×48个热力点，网络分别去预测每个点所属每一类的概率，17代表一个人体的关键点的个数。最后是进行后处理，发现诸如此类的热力图模型的后处理都比较简单，可以参考之前的centernet_tensorrt加速。

3、run-demo

在cpp中运行demo

int main(){
    // cmake使用绝对路径,makefile使用相对路径(包括对应加载图片和模型的.cu文件)
    if(!BaiscTools::build_model("/home/rex/Desktop/tensorrt_learning/trt_cpp/workspace/centernet",1)){
        return -1;
    }
    demoInfer demo;
    //string demo_name = "hrnet"; 
    string demo_name = "hrnet_gpu"; 
    demo.do_infer(demo_name);
    return 0;
}

在python中运行demo

#可能需要调整文件夹的路径
python trtpy/centernet/centernet_inferencv.py

3.1、build_engine

首先将pytorch模型转化为onnx，再将onnx模型幻化为engien，关于转化模型的文章可以参考:https://blog.csdn.net/weixin_42108183/article/details/124969680，拿到onnx模型之后，可以用netron 打开，确认输出没有问题之后，这里的输出的维度为1×17×64×48。下面就可以进行onnx转engine了，依然使用build_engine/batch_1/build_engine_b1.py (此时多batch推理依然没有实现)

3.2、inference

这里只简单展示后处理代码，python

   if use_cupy:
        print("use_cupy")
        t_start = int(time.time() * 1000)
        pred = cp.asarray(pred)
        t_process_start = int(time.time() * 1000)
        print(f"time in numpy2cupy = {t_process_start - t_start} ms")
        pred = pred.reshape(17, 64 * 48)
        indx = cp.argmax(pred, 1)
        scores = cp.max(pred, axis=1)
        keep = scores > 0.2
        indx = indx[keep]
        scores = scores[keep]
        t_end = int(time.time() * 1000)
        print(f'use_cupy postprocess time : {t_end - t_process_start} ms')

    elif use_torch:
        print("use_torch")
        t_start = int(time.time() * 1000)
        pred = torch.from_numpy(pred)
        t_process_start = int(time.time() * 1000)
        print(f"time in numpy2cupy = {t_process_start - t_start} ms")
        pred = pred.reshape(1, 17, 64 * 48).squeeze(0)
        #找到每一行满足条件的点
        keep = pred.max(-1).values > 0.2
        pred = pred[keep]
        scores, indx = torch.max(pred, 1)
        t_end = int(time.time() * 1000)
        print(f'use_torch postprocess time : {t_end - t_process_start} ms')
	#计算坐标
    x = indx % 48
    y = indx / 48
    #恢复至原图大小
    x_os = (x * inv_M[0][0]) + inv_M[0][2]
    y_os = (y * inv_M[1][1]) + inv_M[1][2]
    for x_o, y_o in zip(x_os, y_os):
        cv2.circle(image_o, (int(x_o), int(y_o)), 1, (255, 0, 0), 2)
    cv2.imshow('image', image_o)
    cv2.waitKey(0)

由于最近发现了cupy库，于是尝试用cupy库来进行后处理，结果发现，从numpy中复制数据到cupy中花的时间较长，如果从图像的预处理到推理和后处理都使用cupy的话，应该是能节省一部分时间的，这部分还未实现，总的来说，在这个后处理中，torch花的时间更少，或许cupy还 有其他方法进行加速，毕竟也是刚接触，后续会持续跟进。后处理的方法可以看成从17*3072这个二维矩阵中(每一行代表该类图上出现的以及位置概率)

cpp后处理部分代码

    //预测值
    float* prob = output->cpu<float>();
    float* start = prob;
   	//类别
    int nums = 17;
	//64*48
    int pic_region = 3072;
    for(int i=0;i<output->count();i+=pic_region){
        start = prob + i;
        //位置
        int label = (max_element(start,start+pic_region) - start);
        float score = start[label];
        if(score<0.2)
            continue;
        float x = label % 48 ;
        float y = label / 48 ;
        //这里计算的warpaffine的变换矩阵中的缩放系数 与 pyhton中image_transformer差四倍
        int x_o = (x * d2i[0] * 4) + d2i[2];
        int y_o = (y * d2i[4] * 4) + d2i[5]; 
        cv::circle(img_o, cv::Point((int)x_o,(int)y_o), 1, (0, 0, 255), 2);
        cv::imwrite("hrnet-pred.jpg", img_o);
    }

cpp-cuda后处理代码

static __global__ void hrnet_decode_kernel(
    float* predict,int num_bboxes, int num_classes, float confidence_threshold, 
    float* invert_affine_matrix, float* parray, int max_objects, int NUM_BOX_ELEMENT
){  
    int position = blockDim.x * blockIdx.x + threadIdx.x;
    if (position >= num_bboxes) return;
    // //每隔64*48个位置为 跳到下一行
    float* pitem     = predict + (64*48) * position;
    //pitem 这个类别在图中所有点的概率
    int label = 0;
    //第一个位置的socre
    float* class_confidence = pitem;
    float confidence  = *class_confidence;
    
    for(int i = 1; i < num_classes; ++i, ++class_confidence){
        if(*class_confidence > confidence){
            confidence = *class_confidence;
            label      = i;
        }
    }
    
    if(confidence < confidence_threshold)
        return;
    // printf("%f\n",confidence);
    // printf("%d\n",label-1);
    // printf("%d\n",position);

    int index = atomicAdd(parray, 1);
    if(index >= max_objects)
        return;
    float* pwhxy = pitem;

    float cx  = ((label-1) % 48);
    float cy  = ((label-1) / 48);

    // 第一个位置用来计数
    // left, top, right, bottom, confidence, class, keepflag
    float* pout_item = parray + 1 + index * NUM_BOX_ELEMENT;
    *pout_item++ = cx;
    *pout_item++ = cy;
    *pout_item++ = position;
    *pout_item++ = confidence;
}

其解码方式与cneternet如出一辙，只是减少了nms的操作，最后用一张图来看看hrnet来预测人体关键点的的效果把。

4、总结

本次学习内容主要复习了HRnet的网络结构，其网络的主要特征是多次将不同尺度的特征相融合之后，可以得到比较好的预测效果，同时通过对热力图预测网络的学习，发现其后处理的方式基本一致，而且能够快速的在cuda上进行部署，同时个熟悉了cupy的一些基本的操作。