windows实现libtorch推理行人重识别PCB

**

训练代码

链接: link.

安装libtorch

首先从官网下载libtorch
在这里插入图片描述下载对应cuda版本的libtorch,由于官网只有1.7版本,我采用的是pytorch1.6+libtorch1.6+cuda10.2的组合,这里给出libtorch1.6的下载链接https://download.pytorch.org/libtorch/cu102/libtorch-win-shared-with-deps-1.6.0.zip。下载完解压至相应的文件夹。把libtorch的lib文件夹添加到环境变量里,重启生效。

模型转换

由于python训练保存的是pth文件,是不能在c++中直接调用的,因此需要先转换模型,即转换成c++版本的libtorch可以调用的pt模型,这一步是在python中的pytorch里实现的。
模型的序列化是利用Torch Script来完成的。TorchScript是一种从PyTorch代码创建可序列化和可优化模型的方法。用TorchScript编写的任何代码都可以从Python进程中保存并加载到没有Python依赖关系的进程中。对于一个已经训练好的pytorch模型,官方提供两种方法进行Torch Script的转换:tracing和annotation。
tracing适用于大多数网络,如果你的网络的forward方法中对input有逻辑判断,比如input的size为一个值时走向一个分支,而为另一值时走向另一个分支,那么只能用annotation进行转换。
这里我用的是第一种方法tracing,即需要传一个输入给torch.jit.trace函数,让它输出一次,然后save。模型是采用market1501数据集训练的。示例代码:

from torchvision import models
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from model import PCBModel
from torchvision import datasets, transforms
from PIL import Image
from torch.jit import ScriptModule, script_method, trace
import time
import cv2 as cv


def load_img(img_name):
    loader = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    img = cv.imread(img_name)
    img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
    img = cv.resize(img, (128, 384))
    img = loader(img)
    img = img.unsqueeze(0)
    return img
    
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

model = PCBModel(num_classes=751, num_stripes=6, share_conv=False, return_features=True)

checkpoint = torch.load(r"./net_final.pth")
model.load_state_dict(checkpoint)
model = model.to(device)
model.eval()  # 这句必须要加上,将模型设为eval

# for param in model.parameters():
#         param.requires_grad = False # 如果出现了c++中显存占用比python中很多的情况,转换模型的时候加上这两句话

example = torch.rand(1, 3, 384, 128)
#example = load_img(path) 
example = example.to(device)
#print(example)
print(example.shape)
trace_script_module = torch.jit.trace(model, example)
output = trace_script_module(example)
#print(trace_script_module)
trace_script_module.save(r"./market1501.pt")
print(output)

其中example = load_img(path)的load_img函数是图片预处理函数,必须保证与训练时的图片与处理方式一样,不然模型转换失败,得不到一致的输出,转换成功则会在对应保存的文件夹得到对应的pt模型,这时就可以转移到c++了。

迁移到c++端

VS2017下配置Opencv、libtorch环境,新建一个空项目,选择Release x64,然后和VS配置opencv一样,把LibTorch的include和lib添加到“包含目录”和“库目录”中就行。
在这里插入图片描述
在这里插入图片描述
添加静态库,同样右击项目解决方案选择属性,选择链接器->输入->附加依赖项,添加静态库:
c10.lib
c10_cuda.lib
torch.lib
torch_cpu.lib
torch_cuda.lib
caffe2_nvrtc.lib
caffe2_module_test_dynamic.lib

由于要进行GPU推理,同样右击项目解决方案选择属性,选择链接器->所有选项->附加选项,加入这句话 -INCLUDE:?warp_size@cuda@at@@YAHXZ,不然会出现cuda不可用的问题。

解决提醒std冲突问题,还需修改两个地方:
第一项:属性->C/C++ ->常规->SDL检查->否。
第二项:属性->C/C++ ->语言->符号模式->否。
输入一张自定义图片,预处理、模型推理、后处理等,示例代码:

#include <torch/torch.h>
#include <torch/script.h>

#include <opencv2/opencv.hpp>
#include <vector>

int main()
{

torch::jit::script::Module module = torch::jit::load("pt模型路径", torch::kCUDA); //模型加载到cuda上
std::cout << "模型加载成功" << std::endl;
torch::NoGradGuard no_grad_guard;

cv::Mat image = cv::imread("1.jpg");

//图片预处理,保证跟pytorch训练时的预处理一样的
cv::resize(image, image, cv::Size(128, 384));
cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
//输入数据移到cuda上
torch::Tensor img_tensor = torch::from_blob(image.data, { 1, image.rows, image.cols, 3 }, torch::kByte).to(torch::kCUDA); 
img_tensor = img_tensor.permute({ 0, 3, 1, 2 });
img_tensor = img_tensor.toType(torch::kFloat);
img_tensor = img_tensor.div(255);
img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229);
img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);

std::vector<torch::jit::IValue> inputs;
inputs.push_back(img_tensor);
torch::Tensor output = module.forward(inputs).toTensor();

torch::Tensor features = output.contiguous().view({ 1, -1 });

torch::Tensor fnorm = features.norm(2, 1, false);

torch::Tensor norm_feature = features.div(fnorm.unsqueeze(1)).to(torch::kCPU);

std::vector<float> vec_feature_ouput(norm_feature.data<float>(), norm_feature.data<float>() + norm_feature.numel());
std::cout << "\n" << std::endl;
}

如有错误,欢迎指正,一起交流学习!
未完待续。。。

已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 游动-白 设计师:上身试试 返回首页