利⽤LibTorch调⽤PyTorch训练好的模型
PyTorch如今发布到1.1稳定版本,新增的功能让模型部署变得更为地简单,本⽂记录如何利⽤C++来调⽤PyTorch训练好的模型,其实也是利⽤官⽅强⼤的LibTorch库。
LibTorch的安装
虽然说安装,其实就是下载官⽅的LibTorch包⽽已,从中选择PyTorch(1.1),libtorch,以及cuda的版本,其中会出现下载链接,这⾥为cuda9.0的链接
下载好个路径解压。解压完放在那不动!!
PyTorch模型训练
这⾥我使⽤了最为简单ResNet50的预训练模型,其中保存跟踪模型的代码如下:
import torch
dels as models
from PIL import Image
import numpy as np
image = Image.open("build/airliner.jpg")#图⽚发在了build⽂件夹下
image = size((224,224),Image.ANTIALIAS)
image = np.asarray(image)
image = image /255
image = torch.Tensor(image).unsqueeze_(dim=0)
image = image.permute((0,3,1,2)).float()
model = snet50(pretrained=True)
model = model.eval()cmake如何使用
resnet = ace(model, torch.rand(1,3,224,224))
# output=s(1,3,224,224))
output = resnet(image)
max_index = torch.max(output,1)[1].item()
print(max_index)# ImageNet1000类的类别序号
resnet.save('resnet.pt')
注意这⾥使⽤的是jit中的trace跟踪模型的⽅式,毫⽆疑问最后的输⼊的能够得到飞机的类别,ImageNet1000类的序号类别可以参考
通过该代码能够在根⽬录下产⽣⼀个resnet.pt的⽂件,这个⽂件就是接下来C++所需要调⽤的。
C++调⽤训练好的模型
在写C++调⽤模型的代码之前,先写好CMakeLists⽂件:
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example_torch)
set(CMAKE_PREFIX_PATH "XXX/libtorch")//注意这⾥填⾃⼰解压libtorch时的路径
find_package(Torch REQUIRED)
find_package(OpenCV 3.0 QUIET)
if(NOT OpenCV_FOUND)
find_package(OpenCV 2.4.3 QUIET)
if(NOT OpenCV_FOUND)
message(FATAL_ERROR "OpenCV > 2.4.3 not found.")
endif()
endif()
add_executable(${PROJECT_NAME}"main.cpp")
target_link_libraries(${PROJECT_NAME} ${TORCH_LIBRARIES} ${OpenCV_LIBS})
set_property(TARGET ${PROJECT_NAME} PROPERTY CXX_STANDARD 11)
其中要设置好CMAKE_PREFIX_PATH路径,这个路径就是我们解压libtorch的路径,不然⽆法链接到libtorch库,其中也设置了OpenCV的配置,具体OpenCV的安装这⾥介绍了。
然后就是C++调⽤PyTorch模型的代码
#include<torch/torch.h>
#include<torch/script.h>
#include<iostream>
#include<vector>
#include<opencv2/highgui.hpp>
#include<opencv2/core/core.hpp>
#include<opencv2/opencv.hpp>
void TorchTest(){
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../resnet.pt");
assert(module !=nullptr);
std::cout <<"Load model successful!"<< std::endl;
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::zeros({1,3,224,224}));
at::Tensor output = module->forward(inputs).toTensor();
auto max_result = output.max(1,true);
auto max_index = std::get<1>(max_result).item<float>();
std::cout << max_index << std::endl;
}
void Classfier(cv::Mat &image){
torch::Tensor img_tensor = torch::from_blob(image.data,{1, ws, ls,3}, torch::kByte);
img_tensor = img_tensor.permute({0,3,1,2});
img_tensor = Type(torch::kFloat);
img_tensor = img_tensor.div(255);
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../Train/resnet.pt");
torch::Tensor output = module->forward({img_tensor}).toTensor();
auto max_result = output.max(1,true);
auto max_index = std::get<1>(max_result).item<float>();
std::cout << max_index << std::endl;
}
int main(){
//    TorchTest();
cv::Mat image = cv::imread("airliner.jpg");
cv::resize(image,image, cv::Size(224,224));
std::cout << ws <<" "<< ls <<" "<< image.channels()<< std::endl;
Classfier(image);
return0;
}
其中TorchTest函数只是做了简单的演⽰,⽽Classfier通过OpenCV读取图⽚,并通过libtorch的函数将Mat格式转换成Tensor(注意:这⾥转换了维度,因为OpenCV的维度是[H,W,C], ⽽PyTorch模型需要的是[C,H,W]),最后依然能够输出和Python代码⼀样的答案。
这⾥⽐较重要的⼏个函数有:
torch::from_blob():这个函数将Mat类型转换成Tensor类型。
torch::jit::load():该函数顾名思义就是加载模型的函数。
module->forward():模型前向传播的函数,输⼊值建议使⽤vector类型
max(): 这个函数是libtorch中的max,返回c++中的tuple类型(第⼀个值为维度上最⼤值,第⼆个值为最⼤值的序号)所以使⽤std::get<1> (max_result)来取出序号,这是tuple类型取出⽅式。

版权声明:本站内容均来自互联网,仅供演示用,请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系QQ:729038198,我们将在24小时内删除。