Libtorch前向推理

前言

• 初次使用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;
}