使用LibTorch进行C++调用pytorch模型方式
作者:sq_damowang 时间:2021-09-17 18:56:13
前天由于某些原因需要利用C++调用PyTorch,于是接触到了LibTorch,配了两天最终有了一定的效果,于是记录一下。
环境
PyTorch1.6.0
cuda10.2
opencv4.4.0
VS2017
具体过程
下载LibTorch
去PyTorch官网下载LibTorch包,选择对应的版本,这里我选择Stable(1.6.0),Windows,LibTorch,C++/JAVA,10.2,然后我选择release版本下载,如下图
下载完后先不用管它,之后再用
用pytorch生成模型文件
我先创建了一个python文件,加载resnet50预训练模型,用来生成模型文件,代码如下
import torch
import torchvision.models as models
from PIL import Image
import numpy as np
from torchvision import transforms
model_resnet = models.resnet50(pretrained=True).cuda()
# model_resnet.load_state_dict(torch.load("resnet_Epoch_4_Top1_99.75845336914062.pkl"))
model_resnet.eval()
# 自己选择任意一张图片,并将它的路径写在open方法里,用来读取图像,我这里路径就是‘111.jpg'了
image = Image.open("111.jpg").convert('RGB')
tf = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
])
img = tf(image)
img = img.unsqueeze(dim=0)
print(img.shape)
input = torch.rand(1, 3, 224, 224).cuda()
traced_script_module_resnet = torch.jit.trace(model_resnet, input)
output = traced_script_module_resnet(img.cuda())
print(output.shape)
pred = torch.argmax(output, dim=1)
print(pred)
traced_script_module_resnet.save("model_resnet_jit_cuda.pt")
最后可以生成一个model_resnet_jit_cuda.pt文件,产生的输出如下所示
第一行是我们读取图像的shape,我们读取图片之后经过各种resize,增加维度,把图片数据的shape修改成模型接受的格式,可以看到预测的结果是921,之后我们将用到生成的model_resnet_jit_cuda.pt文件。
VS创建工程并进行环境配置
我在这个python文件路径下创建了这个vs工程Project1
创建完成之后我们打开Project1文件夹,里面内容如下
现在创建VS工程先告一段落,开始进行工程环境配置。把之前下载的LibTorch,解压到当前目录,解压后会出现一个libtorch的文件夹,文件夹目录里的内容为
这里将我框选的文件夹路径配置到工程属性当中,打开刚才新建的VS工程,选择项目为relaese的×64版本
然后点击项目->Project1属性,弹出属性页
在属性页同样注意是release的×64平台,点击VC++目录,在包含目录下加载我之前框出来的include文件夹路径,在库目录下加载框出来的lib文件夹路径,同时,我们也要用到opencv,所以也需要在包含目录下加载opencv的include文件夹与opencv2文件夹,在库目录下加载opencv\build\x64\vc14\lib,如下图
然后在属性页的链接器->输入,添加附加依赖项,首先先把opencv的依赖项添加了
opencv_world440.lib,(如果一直用的Debug模式,就添加opencv_world440d.lib),然后将libtorch/lib里所有后缀为.lib的文件全添加进来,打开这个文件夹
全都写进去,再点击确定,如下图所示
然后点击链接器->命令行,加上/INCLUDE:?warp_size@cuda@at@@YAHXZ 这一句,加上这一句是因为我们要用cuda版本的,如果是cpu版本可以不加。
最后点击C/C++ ->常规的SDL检查,设置为否
点击C/C++ ->语言的符合模式,设置为否
到此我们的配置就全部结束了!最后!复制libtorch/lib文件夹下所有文件,粘贴到工程文件夹Project1/×64/release文件夹里(点击此处的Project1文件夹可以发现里面也有一个×64/release,之前我也纠结是放在哪,然后我都试了一下,发现这个里面是可以不放的)
运行VS2017工程文件
然后我运行VS工程下一个空的main文件,没有报错,配置大致是没问题的,最后添加完整代码,如下
#include <torch/script.h> // One-stop header.
#include <opencv2/opencv.hpp>
#include <iostream>
#include <memory>
//https://pytorch.org/tutorials/advanced/cpp_export.html
std::string image_path = "../../111.jpg";
int main(int argc, const char* argv[]) {
// Deserialize the ScriptModule from a file using torch::jit::load().
//std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("../../model_resnet_jit.pt");
using torch::jit::script::Module;
Module module = torch::jit::load("../../model_resnet_jit_cuda.pt");
module.to(at::kCUDA);
//assert(module != nullptr);
//std::cout << "ok\n";
//输入图像
auto image = cv::imread(image_path, cv::ImreadModes::IMREAD_COLOR);
cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
cv::Mat image_transfomed;
cv::resize(image, image_transfomed, cv::Size(224, 224));
// 转换为Tensor
torch::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.unsqueeze(0);
tensor_image = tensor_image.to(at::kCUDA);
// 网络前向计算
at::Tensor output = module.forward({ tensor_image }).toTensor();
//std::cout << "output:" << output << std::endl;
auto prediction = output.argmax(1);
std::cout << "prediction:" << prediction << std::endl;
int maxk = 3;
auto top3 = std::get<1>(output.topk(maxk, 1, true, true));
std::cout << "top3: " << top3 << '\n';
std::vector<int> res;
for (auto i = 0; i < maxk; i++) {
res.push_back(top3[0][i].item().toInt());
}
for (auto i : res) {
std::cout << i << " ";
}
std::cout << "\n";
system("pause");
}
得到最终输出为921,可以看到和之前的python文件下输出一致,这里还输出了它的top前三,分别是921,787,490。
注意到,我的这两个输出相同的前提条件是:
1、确定加载的是由对应python文件生成的模型!
2、输入的图片是同一张!并且在python下和C++下进行了同样的转换,这里我在python下,将它进行了RGB模型的转换,resize(224, 224),并且将它的每一个元素值除以255.0,转换到0~1之间(ToTensor()方法),最后维度转换为1, 3, 224, 224,在C++中同样需要将BGR模型转化为RGB模型,进行图像缩放至224,224,并且将像素值除以255,将类型转化为float类型,最后维度同样转换为1,3,224,224,再进行网络前向计算。
来源:https://blog.csdn.net/sq_damowang/article/details/109013176