PyTorch模型C++部署

学习一下 专栏收录该内容
3 篇文章 0 订阅

在生产环境中 PyTorch 模型经常需要部署在 C++ 程序中,目前我找的方法有三种:

  • LibTorch: PyTorch 官方 C++ 库
  • ONNX Runtime
  • OpenCV: DNN 模块

示例网络

本文以 PyTorch 官方教程中的 SuperResolution 模型为例,介绍模型在 C++ 中如何部署。模型结构如下:

# Super Resolution model definition in PyTorch
import torch.nn as nn
import torch.nn.init as init
import torch.utils.model_zoo as model_zoo

class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight)

# Create the super-resolution model by using the above model definition.
torch_model = SuperResolutionNet(upscale_factor=3)

然后加载事先训练好的模型权重:

# Load pretrained model weights
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1    # just a random number

# Initialize model with the pretrained weights
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))

# set the model to inference mode
torch_model.eval()

不管是 LibTorch 用的 .pt 格式模型,还是 OpenCV DNN 和 ONNX Runtime 用的 .onnx 模型,都是需要在 Python 中将 PyTorch 模型导出得到的 (不能直接用 .pth 模型)。在模型导出前必须执行 torch_model.eval() 或者 torch_model.train(False) 将模型转为推理模式,因为像 dropout 或 batchnorm 之类的操作在推理模式和训练模式下的行为是不同的。

LibTorch

使用 LibTorch 需要事先将模型导出为 Torch Script (.pt 格式) 文件,参见官方教程。有两种方法:Tracing 和 Scripting。

Tracing 就是提供一个示例输入,让 PyTorch 跑一遍整个网络,将过程中的全部操作记录下来,从而生成 Torch Script 模型:

x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
traced_script_model = torch.jit.trace(torch_model, x)
traced_script_model.save("super_resolution.pt")

Scripting 则是直接分析网络结构转化模型:

traced_script_model = torch.jit.script(torch_model)
traced_script_model.save("super_resolution.pt")

这两种方法各有优缺点:如果模型正向传播的控制流跟输入相关,显然 Tracing 只能得到一种输入下的控制流,此时应该用 Scripting;而当模型使用了一些 Torch Script 不支持的特性,同时模型源码又无法修改时(如果能访问源码,Scripting 可以通过加入注释的方法忽略它们),Scripting 便无能为力了,此时只能考虑 Tracing。另外还有一种说法是 Tracing 得到的模型性能要更好一些。更多关于 Tracing 和 Scripting 的区别可以参考 Mastering TorchScript

然后是 C++ 部分,LibTorch 下载参见 PyTorch: Get Started。CMake 工程中使用 LibTorch 只需要加入 find_package(Torch REQUIRED),并将自己的可执行文件/库链接到 ${TORCH_LIBRARIES} 即可。具体 C++ 代码如下:

#include <opencv2/opencv.hpp>
#include <torch/script.h>

torch::Tensor toTensor(const cv::Mat& image)
{
    // convert 8UC1 image to 4-D float tensor
    CV_Assert(image.type() == CV_8UC1);
    return torch::from_blob(image.data,
        { 1, 1, image.rows, image.cols }, torch::kByte)
        .toType(torch::kFloat32)
        .mul(1.f / 255.f);
}

cv::Mat toMat(const torch::Tensor& tensor)
{
    // convert tensor to 8UC1 image
    using namespace torch;
    Tensor t = tensor.mul(255.f).clip(0, 255).toType(kU8).to(kCPU).squeeze();
    CV_Assert(t.sizes().size() == 2);
    return cv::Mat(t.size(0), t.size(1), CV_8UC1, t.data_ptr()).clone();
}

int main()
{
    const bool use_cuda = false;
    const std::string fn_image = "cat.jpg";
    const std::string fn_model = "super_resolution.pt";

    // load model
    auto module = torch::jit::load(fn_model);
    if (use_cuda)
        module.to(torch::kCUDA);

    // load source image
    auto image = imread(fn_image, cv::IMREAD_GRAYSCALE);
    auto input = toTensor(image);
    if (use_cuda)
        input = input.to(torch::kCUDA);

    // inference
    auto output = module.forward({ input }).toTensor();
    auto result_torch = toMat(output);
    imwrite("result_torch.png", result_torch);
}

ONNX Runtime

ONNX (Open Neural Network Exchange) 是微软和脸书主导的深度学习开发工具生态系统,ONNX Runtime (简称 ORT) 则是微软开发的跨平台高性能机器学习训练与推理加速器,根据官方的说法推理/训练速度最高能有 17X/1.4X 的提升,其优异的性能非常适合深度学习模型部署。

不过想吐槽的是,官方网站的 Linux 预编译包下载指引让人摸不着头脑,建议直接从 Github Release 页面下载。解压后你会发现预编译包里除了一些文档之外,只有头文件和二进制的库文件,没有任何包管理相关 (CMake、pkg-config 之类) 的配置文件。虽然源码 CMakeLists 中明明有 pkg-config 配置文件 .pc 的生成,但不知为何并没有被打包进预编译包。总之,拿到预编译包你没法直接通过 CMake/pkg-config 引入自己的工程。不愧是微软家的,完美继承了 Windows 上项目配置复杂的优良传统。

所以,在 CMake 项目中无法通过 find_package 找到 ONNX Runtime。可以仿照这个仓库,使用 find_pathfind_library 来查找:

find_path(ONNX_RUNTIME_SESSION_INCLUDE_DIRS onnxruntime_cxx_api.h
    HINTS /usr/local/include/onnxruntime/core/session/)
find_path(ONNX_RUNTIME_PROVIDERS_INCLUDE_DIRS cuda_provider_factory.h
    HINTS /usr/local/include/onnxruntime/core/providers/cuda/)
find_library(ONNX_RUNTIME_LIB onnxruntime HINTS /usr/local/lib)

add_executable(inference inference.cpp)
target_include_directories(inference PRIVATE
    ${ONNX_RUNTIME_SESSION_INCLUDE_DIRS}
    ${ONNX_RUNTIME_PROVIDERS_INCLUDE_DIRS})
target_link_libraries(inference PRIVATE ${ONNX_RUNTIME_LIB})

分别指定了包含路径、库路径、链接库,这操作是不是有配置 VS 项目的味道了 🙃

万幸的是,我在 Github Issue 中找到有个人自己写的 CMake 配置文件,在预编译包的根目录下建立 share/cmake/onnxruntime 文件夹,在里面创建 onnxruntimeConfig.cmake 文件,内容为:

# This will define the following variables:
#   onnxruntime_FOUND        -- True if the system has the onnxruntime library
#   onnxruntime_INCLUDE_DIRS -- The include directories for onnxruntime
#   onnxruntime_LIBRARIES    -- Libraries to link against
#   onnxruntime_CXX_FLAGS    -- Additional (required) compiler flags

include(FindPackageHandleStandardArgs)

# Assume we are in <install-prefix>/share/cmake/onnxruntime/onnxruntimeConfig.cmake
get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH)
get_filename_component(onnxruntime_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE)

set(onnxruntime_INCLUDE_DIRS ${onnxruntime_INSTALL_PREFIX}/include)
set(onnxruntime_LIBRARIES onnxruntime)
set(onnxruntime_CXX_FLAGS "") # no flags needed


find_library(onnxruntime_LIBRARY onnxruntime
    PATHS "${onnxruntime_INSTALL_PREFIX}/lib"
)

add_library(onnxruntime SHARED IMPORTED)
set_property(TARGET onnxruntime PROPERTY IMPORTED_LOCATION "${onnxruntime_LIBRARY}")
set_property(TARGET onnxruntime PROPERTY INTERFACE_INCLUDE_DIRECTORIES "${onnxruntime_INCLUDE_DIRS}")
set_property(TARGET onnxruntime PROPERTY INTERFACE_COMPILE_OPTIONS "${onnxruntime_CXX_FLAGS}")

find_package_handle_standard_args(onnxruntime DEFAULT_MSG onnxruntime_LIBRARY onnxruntime_INCLUDE_DIRS)

然后将预编译包的 include、lib、share 三个文件夹拷到系统路径,或者注册用户包,就能在自己 CMake 项目中使用 find_package 找到 onnxruntime 了:

find_package(onnxruntime REQUIRED)
add_executable(t main.cpp)
target_link_libraries(t PRIVATE onnxruntime)

C++ 代码 (主要参考):

#include <cuda_provider_factory.h>
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>

int main()
{
    const bool use_cuda = false;
    const std::string fn_image = "cat.jpg";
    const std::string fn_model = "super_resolution.onnx";

    // environment and options
    Ort::Env env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "SuperResolution");
    Ort::SessionOptions session_options;
    if (use_cuda) {
        // https://github.com/microsoft/onnxruntime/blob/rel-1.6.0/include/onnxruntime/core/providers/cuda/cuda_provider_factory.h#L13
        OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
    }
    session_options.SetGraphOptimizationLevel(
        GraphOptimizationLevel::ORT_ENABLE_ALL);

    // load model and create session
    Ort::Session session(env, fn_model.c_str(), session_options);
    Ort::AllocatorWithDefaultOptions allocator;

    // model info
    const char* input_name = session.GetInputName(0, allocator);
    const char* output_name = session.GetOutputName(0, allocator);
    auto input_dims = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    auto output_dims = session.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    input_dims[0] = output_dims[0] = 1;
    std::vector<const char*> input_names { input_name };
    std::vector<const char*> output_names { output_name };

    // input & output data
    auto image = imread(fn_image, cv::IMREAD_GRAYSCALE);
    CV_Assert(image.type() == CV_8UC1);
    CV_Assert(image.rows == input_dims[2] && image.cols == input_dims[3]);
    cv::Mat blob = cv::dnn::blobFromImage(image, 1.0 / 255.0);
    cv::Mat output(output_dims[2], output_dims[3], CV_32FC1);
    auto memory_info = Ort::MemoryInfo::CreateCpu(
        OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
    std::vector<Ort::Value> input_tensors, output_tensors;
    input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
        memory_info, blob.ptr<float>(), blob.total(), input_dims.data(), input_dims.size()));
    output_tensors.emplace_back(Ort::Value::CreateTensor<float>(
        memory_info, output.ptr<float>(), output.total(), output_dims.data(), output_dims.size()));

    // inference
    session.Run(Ort::RunOptions { nullptr }, input_names.data(), input_tensors.data(), 1,
        output_names.data(), output_tensors.data(), 1);

    cv::Mat result_ort;
    cv::convertScaleAbs(output, result_ort, 255.0);
    cv::imwrite("result_ort.png", result_ort);
}

OpenCV DNN

参见官方文档,OpenCV DNN 模块能够直接加载 onnx 格式的模型,C++ 代码:

#include <opencv2/opencv.hpp>

using namespace cv;
using namespace dnn;

int main()
{
    const bool use_cuda = false;
    const std::string fn_image = "cat.jpg";
    const std::string fn_model = "super_resolution.onnx";

    // load and config model
    Net net = readNetFromONNX(fn_model);
    net.setPreferableBackend(DNN_BACKEND_OPENCV);
    net.setPreferableTarget(DNN_TARGET_CPU);

    // source image
    auto image = imread(fn_image, cv::IMREAD_GRAYSCALE);
    Mat blob;
    blobFromImage(image, blob, 1.0 / 255.0);

    // inference and output
    net.setInput(blob);
    auto output = net.forward();
    int new_size[] = { output.size[2], output.size[3] };
    output = output.reshape(1, 2, new_size);
    convertScaleAbs(output, output, 255.0);
}

OpenCV 支持多种 BackendTarget,但下载的预编译包应该只有 DNN_BACKEND_OPENCV + DNN_TARGET_CPU 的实现,想试其它的需要自己编译。关于不同 Backend 的执行效率可以参考 OpenCV Wiki

效率相关

简单跑了一下 SuperResolution 对比这三种 C++ 部署方法 (测试工程),结果如下:

MethodCPUGPU
LibTorch~31 ms2.5~3.0 ms
OpenCV DNN~39 ms-
ONNX Runtime~21 ms~2 ms

其中,OpenCV 只测了 DNN_BACKEND_OPENCV + DNN_TARGET_CPU 实现 (偷懒没编译 OpenCV)。ONNX Runtime 在 GPU 上运行时,跑第一次耗时约 475 ms,后面下降到约 2 ms,应该是初始化比较耗时吧。

总之,ONNX Runtime 效率上的优势还是很明显的。虽然没有测试 OpenCV 其它 Backend,但估计很难超越 ORT。

  • 1
    点赞
  • 3
    评论
  • 8
    收藏
  • 打赏
    打赏
  • 扫一扫,分享海报

评论 3 您还未登录,请先 登录 后发表或查看评论
©️2022 CSDN 皮肤主题:创作都市 设计师:CSDN官方博客 返回首页

打赏作者

mightbxg

你的鼓励将是我创作的最大动力

¥2 ¥4 ¥6 ¥10 ¥20
输入1-500的整数
余额支付 (余额:-- )
扫码支付
扫码支付:¥2
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值