使用pytorch的C++API部署图像分类模型

使用pytorch的C++API部署图像分类模型

2023年7月10日发(作者:)

使⽤pytorch的C++API部署图像分类模型 在部署pytorch模型时候,使⽤C++ API能有更⾼的效率,本⽂记录使⽤C++ API部署⼀个图像分类模型的过程。1.模型转换 ⾸先需要将pytorch模型转化为Torch Script,Torch Script是PyTorch模型的⼀种表⽰,可以被Torch Script编译器理解,编译和序列化。⽤torch script把torch模型转成c++接⼝可读的模型有两种⽅式:Tracing && Annotation. tracing⽐Annotation简单,但只适合结构固定的⽹络模型,即forward中没有控制流的情况,因为Tracing只会保存运⾏时实际⾛的路径。如果forward函数中有控制流,需要⽤Annotation⽅式实现。本⽂采⽤Tracing的⽅式进⾏模型转换,tracing顾名思义,就是沿着数据运算的路径⾛⼀遍。import torchmodel = ('./weights/best_', map_location="cuda:0")()#

使⽤

⽣成 Module

来跟踪x = (1, 3, 224, 224)x = () # very importanttraced_script_module = (model, x) Module序列化 将ScriptModule序列化后才可以在c++中顺利的读取模型,⽽且在这个过程中不需要任何python依赖。traced_script_("") 得到的 .pt⽂件即转换后的模型⽂件,可以直接在C++环境中使⽤,不⽤依赖于任何python环境。3.c++中加载Script Module 使⽤ torch::jit::load()加载模型。#include // One-stop header.#include #include int main(int argc, const char* argv[]) { if (argc != 2) { std::cerr << "usage: example-app n"; return -1; } // Deserialize the ScriptModule from a file using torch::jit::load(). std::shared_ptr module = torch::jit::load(argv[1]); assert(module != nullptr); std::cout << "okn"; }4.完整的预测⽰例#include #include "torch/script.h"#include "torch/torch.h"#include "opencv2/"#include "opencv2/"#include "opencv2/"#include #include #include #include #include#define CLK_TCK 18.2using namespace std;using namespace cv;int main(){ // load model torch::DeviceType device_type; device_type = torch::kCPU; if (torch::cuda::is_available()) { device_type = torch::kCUDA; } else { device_type = torch::kCPU; } torch::Device device(device_type); torch::jit::script::Module module = torch::jit::load("./"); (device); std::cout<<"load model success"<(getTickCount()); for (int k=0; k<1000; k++){ Mat img = imread("xxxxxx/video_down/wangzhe/"); int img_size = 224; Mat img_resized = (); resize(img, img_resized,Size(img_size, img_size)); Mat img_float; img_tTo(img_float, CV_32F, 1.0f / 255.0f); //归⼀化到[0,1]区间 auto tensor_image = torch::from_blob(img_, {1, img_size, img_size, 3}, torch::kFloat32); //对于⼀张图⽽⾔可使⽤此函数将nhwc格式转换成tensor tensor_image = tensor_e({0, 3, 1, 2});//调整通道顺序,将nhwc转换成nchw tensor_image[0][0] = tensor_image[0][0].sub_(0.485).div_(0.229); tensor_image[0][1] = tensor_image[0][1].sub_(0.456).div_(0.224); tensor_image[0][2] = tensor_image[0][2].sub_(0.406).div_(0.225); tensor_image = tensor_(at::kCUDA); //将tensor放进GPU中处理 torch::Tensor out_tensor = d({tensor_image}).toTensor(); //前向计算 auto results = out_(-1, true); auto softmaxs = std::get<0>(results)[0].softmax(0); auto indexs = std::get<1>(results)[0]; auto idx = indexs[0].item(); string labels[2] = {"normal", "pk"};

string label = labels[idx]; float confidence = softmaxs[0].item() * 100.0f; cout<<"label:"<

设置Opencv的CMake路径set(OpenCV_DIR /usr/local/share/OpenCV)find_package (OpenCV REQUIRED NO_CMAKE_FIND_ROOT_PATH)if(OpenCV_FOUND) INCLUDE_DIRECTORIES(${OpenCV_INCLUDE_DIRS}) message(STATUS "OpenCV library status:") message(STATUS " version: ${OpenCV_VERSION}") message(STATUS " libraries: ${OpenCV_LIBS}") message(STATUS " include path: ${OpenCV_INCLUDE_DIRS}")endif()set(CMAKE_PREFIX_PATH xxx/anaconda3/envs/pytorch2/lib/python3.6/site-packages/torch)find_package(Torch REQUIRED)#设置编译器版本SET(CMAKE_C_COMPILER g++)if(CMAKE_COMPILER_IS_GNUCXX) add_compile_options(-std=c++11 -fno-stack-protector) # very important key in TK1,otherwise will raise an error call stack smashing detected message(STATUS "optional:-std=c++11")endif(CMAKE_COMPILER_IS_GNUCXX)add_executable(${PROJECT_NAME} )target_link_libraries(${PROJECT_NAME} ${TORCH_LIBRARIES} ${OpenCV_LIBS})SET(CMAKE_BUILD_TYPE DEBUG)6. 编译&执⾏mkdir buildcmake ..make或者也可以在QT中打开⽂件,以便于⾃⼰调试、修改代码。特别注意的是⽂件中最后添加⼀句"SET(CMAKE_BUILD_TYPE DEBUG)",才能够在QT中进⾏代码调试。7.测试效果 这⾥使⽤ResNet50做⼀个⼆分类计算1000张图像,使⽤Python API耗时18秒,C++ API耗时15秒,能够节省15%的时间消耗,更利于模型的部署。github:

发布者:admin,转转请注明出处:http://www.yc00.com/xiaochengxu/1688932212a184919.html

相关推荐

发表回复

评论列表(0条)

  • 暂无评论

联系我们

400-800-8888

在线咨询: QQ交谈

邮件:admin@example.com

工作时间:周一至周五,9:30-18:30,节假日休息

关注微信