转自C++部署pytorch模型(一)————将pytorch模型转化为torchscript模型 - 简书,侵删
(一)————将pytorch模型转化为torchscript模型
Pytorch官方提供的C++API名为libtorch,在Pytorch1.0版本之后开始支持windows,以后就可以直接用libtorch来部署Pytorch模型了。
在Windows平台上用C++部署深度学习应用一直很麻烦,caffe使用自定义层还需要自己写backward(),而TensorFlow对Windows平台的支持也不是很好,在Windows上编译有各种坑等着你。而这次Facebook提供的libtorch还是编译好的文件,减少了很多的工作量,估计会吸引一部分人转向pytorch。
下面是我将pytorch的模型转化为libtorch可用的torchscript模型的代码。
import torch
import torchvision
# 加载模型
from models import FaceNetModel
embedding_size = 128
num_classes = 10000
model = FaceNetModel(embedding_size = embedding_size,num_classes = num_classes)
checkpoint = torch.load(r"D:\scripts\simple_facenet\log\checkpoint_epoch1610.pth")
model.load_state_dict(checkpoint['state_dict'])
model.cuda()
model.eval()
# 向模型中输入数据以得到模型参数
example = torch.rand(1,3,182,182).cuda()
traced_script_module = torch.jit.trace(model,example)
# 保存模型
traced_script_module.save("torch_script_eval.pt")
这里有两个地方容易出错
- 如果模型中有dropout或者batchnorm的话,一定要先将模型设置为eval模式,再保存,否则在用libtorch调用后会出现随机干扰;
- example这个张量的尺寸务必与你自己的模型的输入尺寸一直,否则会出现错误。
这个libtorch出来时间不长,还有很多坑等着我去踩,如果有同样在研究pytorch或者libtorch的朋友,可以加个好友,共同学习。
(二)————使用libtorch调用torchscripts模型
上一篇文章讲了如果把pytorch模型转化成torchscripts模型,这篇介绍如何在libtorch中调用
先贴代码,细节等后面再补充
#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
#include <stdlib.h>
#include <string>
#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
int main()
{
// 加载模型
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("D:\\scripts\\simple_facenet\\log\\torch_script_eval.pt");
assert(module != nullptr);
// 转化为GPU模型
module->to(at::kCUDA);
std::cout << "ok\n" << std::endl;
std::vector<torch::jit::IValue> inputs_1,inputs_2;
// 加载图片,这里加载了两张一样的图片
cv::Mat image_1,image_2;
image_1 = cv::imread("D:\\datasets\\lfw_yolo_182\\Aaron_Eckhart\\Aaron_Eckhart_0001.jpg");
image_2 = cv::imread("D:\\datasets\\lfw_yolo_182\\Aaron_Eckhart\\Aaron_Eckhart_0001.jpg");
// 数据转换
// OPENCV 默认格式为BGR,而我训练模型用的是RGB格式,所以需要转换
cv::cvtColor(image_1, image_1, CV_BGR2RGB);
cv::cvtColor(image_2, image_2, CV_BGR2RGB);
cv::Mat img_float_1,img_float_2;
image_1.convertTo(img_float_1, CV_32F, 1.0 / 255);
image_2.convertTo(img_float_2, CV_32F, 1.0 / 255);
cv::resize(img_float_1, img_float_1, cv::Size(182, 182));
cv::resize(img_float_2, img_float_2, cv::Size(182, 182));
auto img_tensor_1 = torch::CPU(torch::kFloat32).tensorFromBlob(img_float_1.data, { 1,182,182,3 });
auto img_tensor_2 = torch::CPU(torch::kFloat32).tensorFromBlob(img_float_2.data, { 1,182,182,3 });
img_tensor_1 = img_tensor_1.permute({ 0,3,1,2 }); // ת��nchw
img_tensor_2 = img_tensor_2.permute({ 0,3,1,2 }); // ת��nchw
auto img_var_1 = torch::autograd::make_variable(img_tensor_1, false);
auto img_var_2 = torch::autograd::make_variable(img_tensor_2, false);
inputs_1.push_back(img_var_1.to(at::kCUDA));
inputs_2.push_back(img_var_2.to(at::kCUDA));
std::cout << "image ready" << std::endl;
// forward
auto output_1 = module->forward(inputs_1).toTensor();
auto output_2 = module->forward(inputs_2).toTensor();
vector<float> feature1(128),feature2(128);
for (int i = 0;i<128;i++)
{
// 转化成Float
feature1[i] = output_1[0][i].item().toFloat();
feature2[i] = output_2[0][i].item().toFloat();
}
/* 计算vector之间的距离
... ...
*/
getchar();
return 0;
}
这里为了方便部署,把最后的结果转化为vector了,对libtorch的API还没有摸得很清楚,不知道是否还有跟高效率的转化方式。
参考:
https://github.com/Jack-An/TorchDemo/blob/master/main.cpp
利用Pytorch的C++前端(libtorch)读取预训练权重并进行预测 - Oldpan的个人博客