前言
- 初次使用Libtorch做前向推理，如有不足请各位指正。
- 在使用Libtorch做retinaface算法前向推理时发现不可以使用python训练时保存的是pth文件，报如下错误:
terminate called after throwing an instance of 'c10::Error'
what(): [enforce fail at inline_container.cc:143] . PytorchStreamReader failed reading zip archive: failed finding central directory
- 参看这篇文章windows实现libtorch推理行人重识别PCB，需要做一次模型转换得到libtorch可以使用的模型，顺利解决问题
使用流程
- 模型转换->读取模型文件->前处理->使用前处理数据创建libtorch张量->libtorch前向推理->获取推理结果->后处理
我的环境
- libtorch:1.4.0+cpu
- Python 3.5.2
- pytorch 1.1.0
模型转换
- 我测试使用的是retinaface算法，参照的是Pytorch_Retinaface 这个工程，先clone这个工程，并下载对应的模型文件
- 在Pytorch_Retinaface创建转换脚本，代码如下：
from torchvision import models
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from models.retinaface import RetinaFace
from torchvision import datasets, transforms
from PIL import Image
from torch.jit import ScriptModule, script_method, trace
import time
import cv2 as cv
cfg_mnet = {
'name': 'mobilenet0.25',
'min_sizes': [[16, 32], [64, 128], [256, 512]],
'steps': [8, 16, 32],
'variance': [0.1, 0.2],
'clip': False,
'loc_weight': 2.0,
'gpu_train': True,
'batch_size': 32,
'ngpu': 1,
'epoch': 250,
'decay1': 190,
'decay2': 220,
'image_size': 640,
'pretrain': True,
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
'in_channel': 32,
'out_channel': 64
}
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device = "cpu"
print(device)
model = RetinaFace(cfg=cfg_mnet, phase = 'test')
checkpoint = torch.load(r"./weights/mobilenet0.25_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, 320, 320)
#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"./retina_face_libtorch_cpu.pt")
print(output)
- 执行上述模型转换脚本，我们便可以得到./retina_face_libtorch_cpu.pt文件
前向推理代码
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include "dtk_retina_face_com.h"
using namespace cv;
using namespace std;
int main(int argc, char** argv)
{
//读取图片
const char* image_path = "../image/lenna.jpg";
cv::Mat image = cv::imread(image_path);
if (image.empty())
{
fprintf(stderr, "cv::imread %s failed\n", image_path);
return -1;
}
//图片前处理
DtkPreInfo_t pre_info;
cv::Mat input_mat = PreProcessing(image, pre_info);
if (input_mat.empty())
{
std::cout << "PreProcessing image is empty!" << std::endl;
return -1;
}
//1.读取模型文件
std::string model_file = "/home/damon/Pytorch_Retinaface/retina_face_libtorch_cpu.pt";
torch::jit::script::Module module = torch::jit::load(model_file);
//2.使用前处理数据创建libtorch张量inputs.
torch::TensorOptions option(torch::kFloat32);
auto img_tensor = torch::from_blob(input_mat.data, { 1, 3, input_mat.rows, input_mat.cols}, option);// opencv H x W x C torch C x H x W
std::vector inputs;
inputs.push_back(img_tensor);
//3.前向推理，返回有多个值不能用toTensor()获取结果，需要用toTuple()获取结果
auto output = module.forward(inputs).toTuple();
//4.获取推理结果
torch::Tensor out0 = output->elements()[0].toTensor();
torch::Tensor out1 = output->elements()[1].toTensor();
torch::Tensor out2 = output->elements()[2].toTensor();
/*
* 5.后处理
*/
return 0;
}
