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
string label = labels[idx]; float confidence = softmaxs[0].item() * 100.0f; cout<<"label:"<
发布者:admin,转转请注明出处:http://www.yc00.com/xiaochengxu/1688932212a184919.html
评论列表(0条)