如何在c++平台上部署pytorch模型?

pytorch接口简单灵活,深受深度学习研究者的喜爱,不少论文及github上的开源代码是用pytorch写的,那么,训练完pytorch模型后,部署到c++平台上,着实让不少人头疼.好在,pytorch开放了libtorch c++接口,至此,caffe, mxnet, tensorflow, pytorch均完善了python以及c++接口,无论是在PC端,还是在移动端,均可满足大多数人的训练和模型部署需求.

深度学习框架趋于大同的情况下,需要耐心对各框架熟悉.鄙人将pytorch c++部署简述如下:

先提一件事情,除了使用libtorch外,确实有很多部署的方法,例如如下:

(1)pytorch模型转为tensorflow模型,利用tensorflow c++等进行部署;

(2)pytorch模型转为onnx模型,利用其他支持onnx的框架进行部署;

(3)pytorch模型转为tensor rt模型,利用tensorrt进行部署等;

诚然,可以,然而不是最好,缺点如下:

(1)有种逃避困难的嫌疑,走捷径固然好,但是不深入;

(2)各种框架间兼容的情况(eg. operator等)并不好,自己的框架自己最清楚;

(3)你给自己挖了另外一个坑.

 

1.注意opencv和libtorch的冲突:

opencv和libtorch联合使用编译时,会报错,类似如下:

error: undefined reference to `cv::imread(std::string const&, int)'

如果你的OpenCV在单独编译使用时没有错误,但是一块编译就出现问题,那么这代表我们的libtorch库和OpenCV库冲突了,冲突原因可能是OpenCV编译OpenCV的C++-ABI版本和libtorch中的不同,所以建议OpenCV最好和libtorch在同样的环境下编译。

2. libtorch和gcc运行时库的冲突

如果你用官网下载的include和libtorch头文件,编译好可执行程序后,run的时候报错,而且你经过调试发现,根本没有进入到main函数里面,那么极可能是和运行时库冲突了,pytorch官网上要求gcc版本最好大于4.9,但是,在实际的部署及生产环境中,这个要求并不合适,详见我的如下博文:

https://blog.csdn.net/seasermy/article/details/99300499

以上两点,均是由于编译情况引起,原因不好查找,隐藏的很深,最好的办法就是,在自己的当前系统环境(eg. gcc4.8.5)下,重新编译libtorch,并进行使用.

3. 模型转换

libtorch不依赖于python,python训练的模型,需要转换为script model才能由libtorch加载,并进行推理.

关于模型转换的例子和帖子很多,如下:

https://github.com/iamhankai/cpp-pytorch/blob/master/tracing.py

https://pytorch.org/docs/master/jit.html // 大多数例子都是摘自官网,还算讲的明白

方法一:tracing

顾名思义,就是随意生成一个输入,走一边推理网络,然后由torch.ji.trace记录一下路径上的信息并保存,这种方法直接简单,容易接受,但缺点是,如果随着输入的不同,网络有多个分支的话,那么会不实用,最直接的一个例子是超分模型,目前的超分辨率模型的研究方向是单一模型实现多倍超分,不仅节省网络占用的空间,而且还容易部署和使用.

如果用tracing的方法转换超分模型的话,那么实现X2,X3,X4的话,会trace出三个模型,不好吧,不好吧,不好吧..

 

方法2:annotation,做标记

顾名思义,给函数加上标识,自己识别多个分支,我们不用具体trace哪个分支,在pytorch1.2中,貌似不加标识就可以,直接torch.jit.script转换为Script module即可.

美中不足是,有些模型写的比较乱,你需要重构一下模型定义的*.py文件,使分支的程序放在第一个毕竟函数里面,否则会报错,这个只能意会,举例如下:

以超分模型中的CARN为例子,参见我的分支:

https://github.com/idealboy/CARN-pytorch/tree/master/carn/model.script // 重构了网络定义

https://github.com/idealboy/CARN-pytorch/tree/master/tools //转换模型的脚本

4. 模型加载及推理

    torch::DeviceType device_type;
    if(torch::cuda::is_available())
    {   
        device_type = torch::kCUDA;
    }   
    else
    {   
        device_type = torch::kCPU;
    }   

    torch::Device device(device_type, dev_id);
    try
    {
        torch::jit::script::Module module = torch::jit::load("model.pt", device);
    }
    catch(const c10::Error &e)
    {
        printf("load error\n");
    }

    // Create a vector of inputs.
    std::string filename = "1.png";
    cv::Mat image = cv::imread(filename.c_str());
    cv::cvtColor(image, image, CV_BGR2RGB);

    cv::Mat img_float;
    image.convertTo(img_float, CV_32F, 1.0 / 255);
    
    {   
        struct timeval start, end;
        gettimeofday(&start, NULL);
        torch::Tensor image_tensor = torch::from_blob(img_float.data, {1,image.rows,image.cols, 3}, torch::kFloat32);
        image_tensor = image_tensor.permute({0,3,1,2});
        image_tensor = image_tensor.toType(torch::kFloat);

        std::vector inputs;
        inputs.push_back(image_tensor.to(device));

        torch::Tensor scale_tensor = torch::tensor({atoi(argv[1])});
        inputs.push_back(scale_tensor);
        torch::Tensor output = module.forward(inputs).toTensor();

        std::cout << scale_tensor << std::endl;
        std::cout << output.sizes() << std::endl;

        output = output.squeeze().detach().permute({1, 2, 0});
        output = output.mul(255).clamp(0, 255).to(torch::kU8);
        output = output.to(torch::kCPU);
        cv::Mat result(output.size(0), output.size(1), CV_8UC3);
        std::memcpy((void *) result.data, output.data_ptr(), sizeof(torch::kU8) * output.numel());
        cv::cvtColor(result, result, CV_RGB2BGR);
        cv::imwrite("sr.jpg", result);
        gettimeofday(&end, NULL);
        float using_time = (end.tv_sec - start.tv_sec) * 1000 + (end.tv_usec - start.tv_usec) / 1000.;
        printf("super resolution using time:%f\n ms", using_time);
  }

 

你可能感兴趣的:(pytorch研究与应用)