pytorch 模型转到c++ torch模型 CenterNet为例

https://pytorch.apachecn.org/docs/1.2/advanced/cpp_export.html

前期：我们拿到的centernet有带DCN版本的，但是CenterNet源码自带编译的，所以先摒弃它，还没尝试如何转到torch里面，使用dlav032模型先。

１使用下面的方法,注意目前只支持较新版本的pytorch，至少0.4是不可以的；

以下代码我是直接写在/src/lib/detectors/base_detector.py  类初始化里面，

self.model.eval()
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 512, 512).cuda()
res_model = torchvision.models.resnet18()
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(self.model,example)
# save model
traced_script_module.save("./model.pt")
    

另外还需要使得模型是单输出单输出，如果输出矩阵宽高相同forward里面concat操作即可，不同的话，我还没想到什么办法，比如下面这个centernet的修改。

    def forward(self, x):
        x = self.base(x)
        x = self.dla_up(x[self.first_level:])
        #x = self.fc(x)
        #y = self.softmax(self.up(x))
        temp_heads = []
        for head in self.heads:
            temp_head = self.__getattr__(head)(x)
            temp_heads.append(temp_head)
        
        re = torch.cat(temp_heads, 1)
        return re

 

c++工程的验证

    int img_size = 512;
    std::shared_ptr<torch::jit::script::Module> simaRpnNet;
    std::string det_model_path = "/data_1/vir/car_detection/train_ws/CenterNet/c++/model.pt";
    torch::NoGradGuard no_grad;
    simaRpnNet = torch::jit::load(det_model_path);
    simaRpnNet->to(at::kCUDA);
    assert(simaRpnNet != nullptr);
    cout << "[INFO] init model done...\n";

    vector<torch::jit::IValue> inputs;  //def an input

    cv::Mat src, image, float_image;
    src = cv::imread("/data_1/vir/car_detection/train_ws/CenterNet/images/00d13b4a-0306-4d08-a740-b1b5c63f94c40.jpg");
    cout<<to_string(src.at<cv::Vec3b>(0, 0)[0])<<endl;
    cout<<to_string(src.at<cv::Vec3b>(0, 0)[1])<<endl;
    cout<<to_string(src.at<cv::Vec3b>(0, 0)[2])<<endl;

    resize(src, image, cv::Size(img_size, img_size), cv::INTER_LINEAR);  // resize 图像
    cout<<to_string(image.at<cv::Vec3b>(0, 0)[0])<<endl;
    cout<<to_string(image.at<cv::Vec3b>(0, 0)[1])<<endl;
    cout<<to_string(image.at<cv::Vec3b>(0, 0)[2])<<endl;
//    cvtColor(image, image, CV_BGR2RGB);  // bgr -> rgb
    image.convertTo(float_image, CV_32F, 1.0 / 255);   //归一化到[0,1]区间 TODO
    float *point_img;
    cout<<float_image.at<cv::Vec3f>(0, 0)[0]<<endl;
    cout<<float_image.at<cv::Vec3f>(0, 0)[1]<<endl;
    cout<<float_image.at<cv::Vec3f>(0, 0)[2]<<endl;

//    point_img = float_image.ptr(32);
//    cout << *(float_image.data) << endl;  //输出一个像素点点值
    auto img_tensor = torch::CPU(torch::kFloat32).tensorFromBlob(float_image.data, {1, img_size, img_size, 3});   //将cv::Mat转成tensor,大小为1,224,224,3
    img_tensor = img_tensor.permute({0, 3, 1, 2});  //调换顺序变为torch输入的格式 1,3,224,224
    //img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229);  //减去均值,除以标准差
    //img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
    //img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);
    auto img_var = torch::autograd::make_variable(img_tensor, false);  //不需要梯度
    inputs.emplace_back(img_var.to(at::kCUDA));  // 把预处理后的图像放入gpu
    torch::Tensor result = simaRpnNet->forward(inputs).toTensor();  //前向传播获取结果
    inputs.pop_back();
    cout<<"Forward over!!!"<<endl;
    for (int i = 0; i < 9; ++i ){
        cout<<"result:"+ std::to_string(i) +" "<< result[0][i][0][0]<<endl;
    }

核对是否正确的时候碰到的几个问题：

1. c++里面输入的时候不做bgr->rgb　python工程里面是做了的。

2. 比较后输出矩阵值完全一样