PyTorch導出JIT模型并用C++ API libtorch調用

本文將介紹如何將一個 PyTorch 模型導出為 JIT 模型并用 PyTorch 的 C++API libtorch運行這個模型。

Step1：導出模型

首先我們進行第一步，用 Python API 來導出模型，由于本文的重點是在后面的部署階段，因此，模型的訓練就不進行了，直接對 torchvision 中自帶的 ResNet50 進行導出。在實際應用中，大家可以對自己訓練好的模型進行導出。

# export_jit_model.py
import torch
import torchvision.models as modelsmodel = models.resnet50(pretrained=True)
model.eval()example_input = torch.rand(1, 3, 224, 224)jit_model = torch.jit.trace(model, example_input)
torch.jit.save(jit_model, 'resnet50_jit.pth')

導出 JIT 模型的方式有兩種：trace 和 script。

我們采用 torch.jit.trace 的方式來導出 JIT 模型，這種方式會根據一個輸入將模型跑一遍，然后記錄下執行過程。這種方式的問題在于對于有分支判斷的模型不能很好的應對，因為一個輸入不能覆蓋到所有的分支。但是在我們 ResNet50 模型中不會遇到分支判斷，因此這里是合適的。關于兩種導出 JIT 模型的方式各自優劣不是本文的中斷，以后會再寫一篇來分析。

在我們的工程目錄 demo 下運行上面的 export_jit_model.py ，會得到一個 JIT 模型件：resnet50_jit.pth。

Step 2：安裝libtorch

接下來我們要安裝 PyTorch 的 C++ API：libtorch。這一步很簡單，直接下載官方預編譯的文件并解壓即可：

wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip
unzip libtorch-shared-with-deps-latest.zip

也解壓在我們的工程目錄 demo 下即可。

Step 3：安裝OpenCV

用 Python 或 C++ 做圖像任務，OpenCV 是經常用到的。如果還沒有安裝的讀者可以參考如下在工程目錄 demo 下進行安裝，構建的過程可能會比較久。已經安裝的讀者可跳過此步驟，一會兒在 CMakeLists.txt 文件中正確地指定本機的 OpenCV 地址即可。

git clone --branch 3.4 --depth 1 https://github.com/opencv/opencv.git
mkdir demo/build && cd demo/build
cmake ..
make -j 6

Step 4：準備測試圖像并用Python測試

我們先準備一張小貓的圖像，并用 PyTorch ResNet50 模型正常跑一下，一會兒與我們 C++ 模型運行的結果對比來驗證 C++ 模型是否被正確的部署。

kitten.jpg：

寫一個腳本用 PyTorch 運行一下模型：

# pytorch_test.pyimport torchvision.models as models
from torchvision.transforms import transforms
import torch
from PIL import Image# normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
all_transforms = transforms.Compose([transforms.Resize(224),transforms.ToTensor()])# normalize])model = models.resnet50(pretrained=True)
model.eval()img = Image.open('kitten.jpg').convert('RGB')
img_tensor = all_transforms(img).unsqueeze(dim=0)
pred = model(img_tensor).squeeze(dim=0)
print(torch.argmax(pred).item())

輸出結果是：282。通過查看ImageNet 1K 類別名與索引的對應關系，可以看到，結果為 tiger cat，模型預測正確。一會兒我們看一下部署后的 C++ 模型是否能正確輸出結果 282。

Step 5：準備cpp源文件

我們下面準備一會要執行的 cpp 源文件，第一次使用 libtorch 的讀者可以先借鑒下面的文件。

這里有幾個點要說一下，不注意可能會犯錯：

1. cv::imread() 默認讀取為三通道BGR，需要進行B/R通道交換，這里采用 cv::cvtColor()實現。

2. 圖像尺寸需要調整到 $224\times 224$，通過 cv::resize() 實現。

3. opencv讀取的圖像矩陣存儲形式：H x W x C, 但是pytorch中 Tensor的存儲為：N x C x H x W, 因此需要進行變換，就是np.transpose()操作，這里使用tensor.permut()實現，效果是一樣的。

4. 數據歸一化，采用 tensor.div(255) 實現。

// test_model.cpp
#include <vector>#include <torch/torch.h>
#include <torch/script.h>#include <opencv2/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>int main(int argc, char* argv[]) {// 加載JIT模型auto module = torch::jit::load(argv[1]);// 加載圖像auto image = cv::imread(argv[2], cv::ImreadModes::IMREAD_COLOR);cv::Mat image_transfomed;cv::resize(image, image_transfomed, cv::Size(224, 224));cv::cvtColor(image_transfomed, image_transfomed, cv::COLOR_BGR2RGB);// 圖像轉換為Tensortorch::Tensor tensor_image = torch::from_blob(image_transfomed.data, {image_transfomed.rows, image_transfomed.cols, 3},torch::kByte);tensor_image = tensor_image.permute({2, 0, 1});// tensor_image = tensor_image.toType(torch::kFloat);tensor_image = tensor_image.div(255.);// tensor_image = tensor_image.sub(0.5);// tensor_image = tensor_image.div(0.5);tensor_image = tensor_image.unsqueeze(0);// 運行模型torch::Tensor output = module.forward({tensor_image}).toTensor();// 結果處理int result = output.argmax().item<int>();std::cout << "The classifiction index is: " << result << std::endl;return 0;
}

Step 6：構建運行驗證

我們先來寫一下 CMakeLists.txt：

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(resnet50)find_package(Torch REQUIRED PATHS ./libtorch)
find_package(OpenCV REQUIRED)add_executable(resnet50  test_model.cpp)
target_link_libraries(resnet50 "${TORCH_LIBRARIES}" "${OpenCV_LIBS}")set_property(TARGET resnet50  PROPERTY CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED TRUE)set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

現在我們的工程目錄 demo 下有以下文件：

CMakeLists.txt  export_jit_model.py  kitten.jpg  libtorch  pytorch_test.py  resnet50_jit.pth  test_model.cpp

然后開始用 CMake 構建工程：

mkdir build && cd build
OpenCV_DIR=[YOUR_PATH_TO_OPENCV]/opencv/build cmake ..
make

整個過程沒有報錯的話我們就已經構建完成了，會得到一個可執行文件 resnet50 在工程目錄 demo 下。

接下來我們執行，并驗證運行結果是否與 PyTorch 的結果一致：

./build/resnet50 resnet50_jit.pth kitten.jpg

輸出：

The classifiction index is: 282

運行成功并且結果正確。

Ref：

https://www.jianshu.com/p/7cddc09ca7a4

https://blog.csdn.net/cxx654/article/details/115916275

https://zhuanlan.zhihu.com/p/370455320