PyTorch導(dǎo)出JIT模型并用C++ API libtorch調(diào)用

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

Step1：導(dǎo)出模型

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

# 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')

導(dǎo)出 JIT 模型的方式有兩種：trace 和 script。

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

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

Step 2：安裝libtorch

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

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++ 做圖像任務(wù)，OpenCV 是經(jīng)常用到的。如果還沒有安裝的讀者可以參考如下在工程目錄 demo 下進行安裝，構(gòu)建的過程可能會比較久。已經(jīng)安裝的讀者可跳過此步驟，一會兒在 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++ 模型運行的結(jié)果對比來驗證 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())

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

Step 5：準備cpp源文件

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

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

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

圖像尺寸需要調(diào)整到 $224\times 224$，通過 cv::resize() 實現(xiàn)。

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

數(shù)據(jù)歸一化，采用 tensor.div(255) 實現(xiàn)。

// 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);// 圖像轉(zhuǎn)換為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();// 結(jié)果處理int result = output.argmax().item<int>();std::cout << "The classifiction index is: " << result << std::endl;return 0; }

Step 6：構(gòu)建運行驗證

我們先來寫一下 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}")

現(xiàn)在我們的工程目錄 demo 下有以下文件：

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

然后開始用 CMake 構(gòu)建工程：

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

整個過程沒有報錯的話我們就已經(jīng)構(gòu)建完成了，會得到一個可執(zhí)行文件 resnet50 在工程目錄 demo 下。

接下來我們執(zhí)行，并驗證運行結(jié)果是否與 PyTorch 的結(jié)果一致：

./build/resnet50 resnet50_jit.pth kitten.jpg

輸出：

The classifiction index is: 282

運行成功并且結(jié)果正確。

Ref：

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

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

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

總結(jié)

以上是生活随笔為你收集整理的PyTorch导出JIT模型并用C++ API libtorch调用的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。