【CV学习笔记】pytorch2onnx2trt

1、前言

之前在工作中遇到pytorch模型转tensorrt模型的时候，面对不同的代码仓库，有不同的转化方法，后来在学习过程中见到了比较统一的onnx模型，可以将pytorch模型转化为onnx模型，然后再由onnx-tensorrt将onnx转化为engine，本次学习内容将对pytorch2onnx2trt的方法进行系统性的学习，并记录下来。

代码地址:https://github.com/Rex-LK/tensorrt_learning

欢迎正在学习或者想学的CV的同学进群一起讨论与学习，v:Rex1586662742，q群:468713665

2、pytorch2onnx

下载好项目后，打开项目，以代码中的onnx/demo/centernet为例，其中predict.py可对单张图片进行检测,同时也可以将pytorch模型转成onnx模型，只要修改predict.py中的参数 mode = "export_onnx"即可，执行完毕后，onnx模型为model_data/models.onnx使用netron 在浏览器里面查看该onnx模型是否有问题。

netron model_data/models.onnx

拉到最下面，发现最下面有三个分支，刚好对应centerner网络结构中最后的三个特征层，分别用来预测物体的类别、宽高、中心。但从图中可以看出，最后的output只，拿到了预测的类别分数，转成engine后，右边两个将会被抛弃，这不是我们想要看到的情形。

同时，一些繁琐的后处理代码可以在python中进行，尽量避免在c++出现。首先查看原项目的后处理方法，

 outputs = decode_bbox(outputs[0], outputs[1], outputs[2], self.confidence, self.cuda)
 results = postprocess(outputs, self.nms, image_shape, self.input_shape, self.letterbox_image, self.nms_iou)

在源代码发先者两个后处理步骤，可以先将这两个方法先扔到onnx中，如果export_onnx出现错误，则说明里面的某些操作暂时不支持转onnx，于是需要进行拆分，下面是我尝试多次得到的结果，在mypredict.py中

    def pool_nms(self, heat, kernel=3):
        pad = (kernel - 1) // 2
        hmax = nn.functional.max_pool2d(heat, (kernel, kernel), stride=1, padding=pad)
        keep = (hmax == heat).float()
        return heat * keep

    def forward(self, x):
        x = self.model(x)
        #热力图
        hms = self.pool_nms(x[0]).view(1, self.num_class, self.feature_map * self.feature_map).permute(0, 2, 1)
        
        whs = x[1].view(1, 2, self.feature_map * self.feature_map).permute(0, 2, 1)
        xys = x[2].view(1, 2, self.feature_map * self.feature_map).permute(0, 2, 1)
        y = torch.cat((hms,whs),dim=2)
        y = torch.cat((y,xys),dim=2)
        return y

这里可以看出在onnx中，帮助我们实现了一个最大池化实现NMS的过程，然后将热力图hm、宽高wh、中心xy拼接在一个tensor里面，方便在c++中进行后处理。利用export_onnx.py再次生成onnx模型，利用netron打开,拉到最下面

发现最后的分支全部连接在一起了，最后的1×16384×24表示在图中预测128*128个目标，每个目标的类别数为20，宽高为2，中心为2，有了这个onnx模型之后，就可以利用tensorrt模型进行加速了。

3、onnx2trt

首相将生成好的onnx模型放在/workspace下面，如/workspce/centernet.onnx，可以使用makefile和cmake两种方式进行编译，使用cmake方式时，建议模型和图片使用绝对路径，下面以cmake方法为例,打开src/main.py文件，里面包含了构造引擎和推理。

    //构造引擎
    if(!BaiscTools::build_model("/home/rex/Desktop/deeplearning_rex/onnx2trt/workspace/centernet",1)){
        return -1;
    }
	//推理demo
    demoInfer demo;
    string demo_name = "centernet";
    demo.do_infer(demo_name);
    return 0;
}

修改demo-infer.cpp中的图片路径和模型路径(全局)，然后执行

mkdir build && cd build
cmake .. && make -j

在workspace下面 生成了 demo_infer ,

执行./demo_infer 即可生成引擎，然后进行推理，推理结果在build文件夹下

可以看出，预测效果还是很不错的，下面将介绍centernet在c++中的后处理，在demo-iner.cpp文件中，找到centernet_inference函数。大致可分为三个部分，前面是图像预处理，中间的模型推理，最后是后处理。首先可以参考Python的代码，将confidence大于阈值的目标提取出来。

	//预测值
	float* prob = output->cpu<float>();
    int num = 20;  //类别
    float *start = prob;
    int count = output->count(); 
     vector<bbox> boxes;
    for(int i=0;i<count;i+=24){
        //现在有128*128个点  就有128*128行，每行24个，前20个为类别，后四个为 w,h,x,y
        start = prob+i;
        //找到得分最大索引
        int label = max_element(start,start+num) - start;
        float confidence = start[label];
        if(confidence<0.3)
            continue;
        float w = start[ num];
        float h = start[num+1];
        //这里的x,y 还要加上偏移量
        float x = start[num+2] + (i/24)%128;
        float y = start[num+3] + (i/24)/128;
        float left   = (x - w * 0.5) /128 * img_h;
        float top    = (y - h * 0.5) /128 * img_w;
        float right  = (x + w * 0.5) /128 * img_h;
        float bottom = (y + h * 0.5) /128 * img_w;
        //将每个框放在boxes中
        boxes.emplace_back(left, top, right, bottom, confidence, (float)label);
    }

通过上面的函数得到了confidence大于阈值的框，接下来就要进行NMS操作了。

    //nms
    sort(boxes.begin(), boxes.end(), [](bbox& a, bbox& b){return a.confidence > b.confidence;});
    //然后使用vector<bool>来标记是是否要删除box
    vector<bool> remove_flags(boxes.size());
    vector<bbox> box_result;
    box_result.reserve(boxes.size());

    //计算两个box之间的IOU
    auto iou = [](const bbox& a, const bbox& b){
        float cross_left   = max(a.left, b.left);
        float cross_top    = max(a.top, b.top);
        float cross_right  = min(a.right, b.right);
        float cross_bottom = min(a.bottom, b.bottom);
        //计算重叠部分的面积,注意面积非0
        float cross_area = max(0.0f, cross_right - cross_left) * max(0.0f, cross_bottom - cross_top);
        //A面积+B面积 - 一个重叠面积
        float union_area = max(0.0f, a.right - a.left) * max(0.0f, a.bottom - a.top) 
                         + max(0.0f, b.right - b.left) * max(0.0f, b.bottom - b.top) - cross_area;
        if(cross_area == 0 || union_area == 0) return 0.0f;
        return cross_area/union_area;
    };

    for(int i = 0;i<boxes.size();i++){
        if(remove_flags[i]) continue;
        auto &ibox  = boxes[i];
        //第一个box必定不会remove
        box_result.emplace_back(ibox);
        //第一个box与 之后的box两两比较
        for(int j = i + 1;j<boxes.size();j++){
            if(remove_flags[j]) continue;
            auto& jbox = boxes[j];
            //如果两个box的lable相同
            if(ibox.label == jbox.label){
                //则比较IOU
                if(iou(ibox, jbox) >= 0.3)
                    remove_flags[j] = true;
            }
        }
    }
    //画框
    for(auto& box : box_result){
        cv::rectangle(img_o, cv::Point(box.left, box.top), cv::Point(box.right, box.bottom), cv::Scalar(0, 255, 0), 2);
        cv::putText(img_o, cv::format("%.2f", box.confidence), cv::Point(box.left, box.top - 7), 0, 0.8, cv::Scalar(0, 0, 255), 2, 16);
    }
    cv::imwrite("centernet-pred-street.jpg", img_o);

4、总结

本次学习onnx转tensorrt的方式，学习了如何从torch转onnx，然后再从onnx转tensorrt，对c++的后处理方法得到了进一步的了解，主要是参考python的代码来指导c++代码，后续会持续更新其他模型