亲宝软件园·资讯

展开

使用LibTorch进行C++调用pytorch模型方式

sq_damowang 人气:0

前天由于某些原因需要利用C++调用PyTorch,于是接触到了LibTorch,配了两天最终有了一定的效果,于是记录一下。

环境

具体过程

下载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>

//http://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,再进行网络前向计算。

总结

以上为个人经验,希望能给大家一个参考,也希望大家多多支持。

加载全部内容

相关教程
猜你喜欢
用户评论