北京网站编程培训,wordpress胖子马,东华软件是外包公司吗,山东建设厅网站PyTorch导出JIT模型并用C API libtorch调用
本文将介绍如何将一个 PyTorch 模型导出为 JIT 模型并用 PyTorch 的 CAPI libtorch运行这个模型。
Step1#xff1a;导出模型
首先我们进行第一步#xff0c;用 Python API 来导出模型#xff0c;由于本文的重点是在后面的部署…PyTorch导出JIT模型并用C API libtorch调用
本文将介绍如何将一个 PyTorch 模型导出为 JIT 模型并用 PyTorch 的 CAPI libtorch运行这个模型。
Step1导出模型
首先我们进行第一步用 Python API 来导出模型由于本文的重点是在后面的部署阶段因此模型的训练就不进行了直接对 torchvision 中自带的 ResNet50 进行导出。在实际应用中大家可以对自己训练好的模型进行导出。
# export_jit_model.py
import torch
import torchvision.models as modelsmodel models.resnet50(pretrainedTrue)
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 APIlibtorch。这一步很简单直接下载官方预编译的文件并解压即可
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 6Step 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(pretrainedTrue)
model.eval()img Image.open(kitten.jpg).convert(RGB)
img_tensor all_transforms(img).unsqueeze(dim0)
pred model(img_tensor).squeeze(dim0)
print(torch.argmax(pred).item())输出结果是282。通过查看ImageNet 1K 类别名与索引的对应关系可以看到结果为 tiger cat模型预测正确。一会儿我们看一下部署后的 C 模型是否能正确输出结果 282。
Step 5准备cpp源文件
我们下面准备一会要执行的 cpp 源文件第一次使用 libtorch 的读者可以先借鉴下面的文件。
这里有几个点要说一下不注意可能会犯错 cv::imread() 默认读取为三通道BGR需要进行B/R通道交换这里采用 cv::cvtColor()实现。 图像尺寸需要调整到 224×224224\times 224224×224通过 cv::resize() 实现。 opencv读取的图像矩阵存储形式H x W x C, 但是pytorch中 Tensor的存储为N x C x H x W, 因此需要进行变换就是np.transpose()操作这里使用tensor.permut()实现效果是一样的。 数据归一化采用 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.hppint 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().itemint();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