PyTorch模型C++部署
在生产环境中 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 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_path 和 find_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 /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 OpenCV DNN
参见官方文档，OpenCV DNN 模块能够直接加载 onnx 格式的模型，C++ 代码：
#include OpenCV 支持多种 Backend 和 Target，但下载的预编译包应该只有 DNN_BACKEND_OPENCV + DNN_TARGET_CPU 的实现，想试其它的需要自己编译。关于不同 Backend 的执行效率可以参考 OpenCV Wiki。
效率相关
简单跑了一下 SuperResolution 对比这三种 C++ 部署方法 (测试工程)，结果如下：
| Method | CPU | GPU |
|---|---|---|
| LibTorch | ~31 ms | 2.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。