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

目录

• 环境
• 具体过程
• 下载LibTorch
• 用pytorch生成模型文件
• VS创建工程并进行环境配置
• 运行VS2017工程文件
• 总结

前天由于某些原因需要利用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，再进行网络前向计算。

总结

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