pytorch 模型轉到c++　torch模型　CenterNet爲例

前期：我們拿到的centernet有帶DCN版本的，但是CenterNet源碼自帶編譯的，所以先摒棄它，還沒嘗試如何轉到torch裏面，使用dlav032模型先。

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

# 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. 比較後輸出矩陣值完全一樣