利用LibTorch部署PyTorch模型

川长思鸟来 2022-09-16 13:28 237阅读 0赞

PyTorch如今发布到1.1稳定版本，新增的功能让模型部署变得更为地简单，本文记录如何利用C++来调用PyTorch训练好的模型，其实也是利用官方强大的LibTorch库。

## LibTorch的安装 ##

虽然说安装，其实就是下载官方的LibTorch包而已，从官方网站[PyTorch][]中选择PyTorch(1.1)，libtorch，以及cuda的版本，其中会出现下载链接，这里为cuda9.0的链接

[https://download.pytorch.org/libtorch/cu90/libtorch-shared-with-deps-latest.zip][https_download.pytorch.org_libtorch_cu90_libtorch-shared-with-deps-latest.zip]

下载好找个路径解压。解压完放在那不动！！!

## PyTorch模型训练 ##

这里我使用了最为简单ResNet50的预训练模型，其中保存跟踪模型的代码如下：

import torch
    import torchvision.models as models
    from PIL import Image
    import numpy as np
    image = Image.open("build/airliner.jpg") #图片发在了build文件夹下
    image = image.resize((224, 224),Image.ANTIALIAS)
    image = np.asarray(image)
    image = image / 255
    image = torch.Tensor(image).unsqueeze_(dim=0)
    image = image.permute((0, 3, 1, 2)).float()
    
    model = models.resnet50(pretrained=True)
    model = model.eval()
    resnet = torch.jit.trace(model, torch.rand(1,3,224,224))
    # output=resnet(torch.ones(1,3,224,224))
    output = resnet(image)
    max_index = torch.max(output, 1)[1].item()
    print(max_index) # ImageNet1000类的类别序
    resnet.save('resnet.pt')

注意这里使用的是 jit 中的trace跟踪模型的方式，毫无疑问最后的输入的能够得到飞机的类别，ImageNet1000类的序号类别可以参考[此处][Link 1]。

通过该代码能够在根目录下产生一个 resnet.pt 的文件，这个文件就是接下来C++所需要调用的。

## C++调用训练好的模型 ##

在写C++调用模型的代码之前，先写好CMakeLists文件：

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
    
    project(example_torch)
    set(CMAKE_PREFIX_PATH "XXX/libtorch") //注意这里填自己解压libtorch时的路径
    
    find_package(Torch REQUIRED)
    find_package(OpenCV 3.0 QUIET)
    if(NOT OpenCV_FOUND)
        find_package(OpenCV 2.4.3 QUIET)
        if(NOT OpenCV_FOUND)
            message(FATAL_ERROR "OpenCV > 2.4.3 not found.")
        endif()
    endif()
    add_executable(${
        PROJECT_NAME} "main.cpp")
    target_link_libraries(${
        PROJECT_NAME} ${
        TORCH_LIBRARIES} ${
        OpenCV_LIBS})
    set_property(TARGET ${
        PROJECT_NAME} PROPERTY CXX_STANDARD 11)```

其中要设置好CMAKE\_PREFIX\_PATH路径，这个路径就是我们解压libtorch的路径，不然无法链接到libtorch库,其中也设置了OpenCV的配置，具体OpenCV的安装这里介绍了。

然后就是C++调用PyTorch模型的代码

#include <torch/torch.h>
    #include <torch/script.h>
    #include <iostream>
    #include <vector>
    #include <opencv2/highgui.hpp>
    #include <opencv2/core/core.hpp>
    #include <opencv2/opencv.hpp>
    
    void TorchTest(){
        
        std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../resnet.pt");
        assert(module != nullptr);
        std::cout << "Load model successful!" << std::endl;
        std::vector<torch::jit::IValue> inputs;
        inputs.push_back(torch::zeros({
        1,3,224,224}));
        at::Tensor output = module->forward(inputs).toTensor();
        auto max_result = output.max(1, true);
        auto max_index = std::get<1>(max_result).item<float>();
        std::cout << max_index << std::endl;
    }
    
    void Classfier(cv::Mat &image){
        
        torch::Tensor img_tensor = torch::from_blob(image.data, {
        1, image.rows, image.cols, 3}, torch::kByte);
        img_tensor = img_tensor.permute({
        0, 3, 1, 2});
        img_tensor = img_tensor.toType(torch::kFloat);
        img_tensor = img_tensor.div(255);
        std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../Train/resnet.pt");
        torch::Tensor output = module->forward({
        img_tensor}).toTensor();
        auto max_result = output.max(1, true);
        auto max_index = std::get<1>(max_result).item<float>();
        std::cout << max_index << std::endl;
    
    }
    
    int main() {
        
    //    TorchTest();
        cv::Mat image = cv::imread("airliner.jpg");
        cv::resize(image,image, cv::Size(224,224));
        std::cout << image.rows <<" " << image.cols <<" " << image.channels() << std::endl;
        Classfier(image);
        return 0;
    }

其中**TorchTest**函数只是做了简单的演示，而**Classfier**通过OpenCV读取图片，并通过libtorch的函数将**Mat**格式转换成**Tensor**(注意：这里转换了维度，因为OpenCV的维度是\[H,W,C\], 而PyTorch模型需要的是\[C,H,W\])，最后依然能够输出和Python代码一样的答案。

这里比较重要的几个函数有：

**torch::from\_blob()**: 这个函数将`Mat`类型转换成`Tensor`类型。

**torch::jit::load()**: 该函数顾名思义就是加载模型的函数。

**module->forward()**: 模型前向传播的函数，输入值建议使用vector类型

**max()**: 这个函数是libtorch中的max，返回c++中的tuple类型（第一个值为维度上最大值，第二个值为最大值的序号）所以使用std::get<1>(max\_result)来取出序号，这是tuple类型取出方式。

[PyTorch]: https://link.zhihu.com/?target=https://pytorch.org/
[https_download.pytorch.org_libtorch_cu90_libtorch-shared-with-deps-latest.zip]: https://download.pytorch.org/libtorch/cu90/libtorch-shared-with-deps-latest.zip
[Link 1]: https://blog.csdn.net/weixin_41770169/article/details/80482942