LibTorch Export

导出torchscript模型核心代码 modnet示例可能的错误构建C++程序安装libtorch 安装opencv CMakeLists.txt exp.cpp 命令 Windows上配置环境配置 CMakeLists.txt c++

导出torchscript模型

核心代码


model = Net()
model.to(device)
model.eval()
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, output_path)

modnet示例


"""
Export TorchScript model of MODNet

Arguments:
    --ckpt-path: path of the checkpoint that will be converted
    --output-path: path for saving the TorchScript model

Example:
    python export_torchscript.py \\
        --ckpt-path=modnet_photographic_portrait_matting.ckpt \\
        --output-path=modnet_photographic_portrait_matting.torchscript
"""

import os
import argparse

import torch
import torch.nn as nn
import torch.nn.functional as F

from . import modnet_torchscript


if __name__ == '__main__':
    # define cmd arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('--ckpt-path', type=str, required=True, help='path of the checkpoint that will be converted')
    parser.add_argument('--output-path', type=str, required=True, help='path for saving the TorchScript model')
    args = parser.parse_args()

    # check input arguments

    if not os.path.exists(args.ckpt_path):
        print(args.ckpt_path)
        print('Cannot find checkpoint path: {0}'.format(args.ckpt_path))
        exit()

    # create MODNet and load the pre-trained ckpt

    modnet = modnet_torchscript.MODNet(backbone_pretrained=False)

    #modnet = nn.DataParallel(modnet).cuda()
    #state_dict = torch.load(args.ckpt_path)
    #modnet.load_state_dict(state_dict)

    state_dict = torch.load(args.ckpt_path, map_location='cpu')
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]  # remove module.
        new_state_dict[name] = v
    modnet.load_state_dict(new_state_dict)
    modnet.to('cpu')
    modnet.eval()

    # export to TorchScript model

    scripted_model = torch.jit.script(modnet)
    torch.jit.save(scripted_model, os.path.join(args.output_path))

可能的错误

torchscript对网络模型定义文件的写法要求非常严格，可能报不能用下标访问module，不能用lambda表达式等等的错误，这时候要么换一种方式写网络定义的语句，要么升级一下torch版本，高版本的对语法更宽容

导出分CPU、GPU版本，需要注意

构建C++程序

安装libtorch

只需要到官网下载适配linux c++版本的libtorch，具体版本可以修改下载地址的一些关键字符，解压之后即可用

安装opencv

具体见Linux opencv安装方法的笔记，opencv官网上也有具体步骤

CMakeLists.txt



cmake_minimum_required(VERSION 3.0 FATAL_ERROR) #cmake版本要求
project(exp)  #exp代表项目名
SET(CMAKE_C_COMPILER g++) #设定编译器
add_definitions(--std=c++14) # 定义c++标准 #注意自己的c++版本， 根据电脑配置定，也可能是-std=c++11

# 指定libTorch位置 这样不用在cmake指令后面写prefix

set(Torch_DIR /home/cz/workspace/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED) #自动地在Torch_DIR里找相关依赖

# 打印一下Torch的相关变量，看是否正确读入

message(STATUS "Torch library status:")
message(STATUS "    version: ${TORCH_VERSION}")
message(STATUS "    libraries: ${TORCH_LIBS}")
message(STATUS "    include path: ${TORCH_INCLUDE_DIRS}")
message(STATUS "    torch lib : ${TORCH_LIBRARIES} ")

# 需确保安装好OpenCV

find_package(OpenCV REQUIRED) #这里没有SET opencv的lib路径是因为opencv是已经写到环境变量里面去了
message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")
message(STATUS "    torch lib : ${TORCH_LIBRARIES} ")

# 包含头文件include
# include_directories(${OpenCV_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS})

add_executable(exp matting.cpp) # 编译后的可执行文件是exp
target_link_libraries(exp ${TORCH_LIBRARIES} ${OpenCV_LIBS}) # 链接

exp.cpp


#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgcodecs.hpp"
#include <vector>
#include <chrono>
#include <string>
#include <vector>
using namespace cv;
using namespace std;


torch::Tensor process( cv::Mat& image, torch::Device device, int img_W, int img_H)
{
    vector <float> mean_ = {0.5, 0.5, 0.5}; #均值
    vector <float> std_ = {0.5, 0.5, 0.5}; #方差
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);// bgr -> rgb
    cv::Mat img_float;
//    image.convertTo(img_float, CV_32F, 1.0 / 255);//归一化到[0,1]区间,
    cv::resize(image, img_float, cv::Size(img_W, img_H)); # resize

    std::vector<int64_t> dims = {1, img_H, img_W, 3}; #从Mat
    torch::Tensor img_var = torch::from_blob(img_float.data, dims, torch::kByte).to(device);//将图像矩阵转化成指定维度张量
    img_var = img_var.permute({0,3,1,2});//将张量的参数顺序转化为 torch输入的格式 NCHW
    img_var = img_var.toType(torch::kFloat); //转换浮点数类型
    img_var = img_var.div(255); // 收缩到0-1

    //对各通道做归一化操作
    for (int i = 0; i < 3; i++) {
        img_var[0][i] = img_var[0][i].sub_(mean_[i]).div_(std_[i]);
    }
    //关闭自动微分
    img_var = torch::autograd::make_variable(img_var, false);
    return img_var;
}

int main(){


    string img_path = "../person.jpg"; //测试图片路径
    string model_path = "../model.pt"; //torchscript模型路径
    int img_W = 640;
    int img_H = 480;

    //加载模型到CPU
    torch::DeviceType device_type;
    device_type = torch::kCPU;
    torch::Device device(device_type);
    torch::jit::script::Module module = torch::jit::load(model_path);//
    module.to(torch::kCPU);
    module.eval();
    std::cout<<"Model loading complete"<<std::endl;

     // 读取图片
    cv::Mat image = cv::imread(img_path, cv::ImreadModes::IMREAD_COLOR);
    if (image.empty())
        std::cout<<"Can not load image!"<<std::endl;
    // 预处理 转tensor
    torch::Tensor img_var = process(image, device, img_W, img_H);

    //forward 一般是以一个vector<tensor>传进去
    torch::Tensor result = module.forward({img_var}).toTensor();
    std::cout<<"Inference complete!"<<std::endl;
    //后处理
    result = result.squeeze();//结果是(1, 1, H, W)删除1的维度
    result = result.mul(255).to(torch::kU8) ; //把alpha转到0-255
    cv::Mat pts_mat(cv::Size(img_W, img_H), CV_8U, result.data_ptr()); // tensor转Mat
    cv::imwrite("../matting_person.png", pts_mat); //存图
    std::cout<<"matting complete!"<<std::endl;
	return 0;
}

命令


mkdir build
cd build
cmake ..
make

Windows上配置

环境配置

libtorch下载对应版本，libtorch\lib配置环境变量，因为运行时需要torch_cpu.dll等动态链接库，或者将这些dll拷贝到build后的工程目录

opencv安装，build\x64\vc14\bin build\x64\vc15\bin配置环境变量

CMakeLists.txt

与linux上大同小异，不同的是它会忽略掉一些语句比如C14、编译器选项等，默认采用系统带的MVC编译器



cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(exp)  #exp代表项目名

SET(CMAKE_C_COMPILER g++)
add_definitions(--std=c++14) #这两条语句被忽略了

# 指定libTorch位置
set(Torch_DIR C:/Users/cz/Desktop/libtorch-win-shared-with-deps-1.6.0+cpu/libtorch/share/cmake/Torch)
find_package(Torch REQUIRED)
message(STATUS "Torch library status:")
message(STATUS "    version: ${TORCH_VERSION}")
message(STATUS "    libraries: ${TORCH_LIBS}")
message(STATUS "    include path: ${TORCH_INCLUDE_DIRS}")
message(STATUS "    torch lib : ${TORCH_LIBRARIES} ")

# 需确保安装好OpenCV
set(OpenCV_DIR E:/opencv-3.4.8/opencv/build)
find_package(OpenCV REQUIRED)
message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")
message(STATUS "    torch lib : ${TORCH_LIBRARIES} ")

# 包含头文件include
# include_directories(${OpenCV_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS})

add_executable(exp matting.cpp)
target_link_libraries(exp ${TORCH_LIBRARIES} ${OpenCV_LIBS})

c++


#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
#include "opencv2/core.hpp"
#include "opencv2/imgproc.hpp"
#include "opencv2/highgui.hpp"
#include "opencv2/imgcodecs.hpp"
#include <vector>
#include <chrono>
#include <string>
#include <vector>
#include <time.h>
//using namespace cv;
//using namespace std;


torch::Tensor process( cv::Mat& image, torch::Device device, int img_W, int img_H)
{
    std::vector <float> mean_ = {0.5, 0.5, 0.5};
    std::vector <float> std_ = {0.5, 0.5, 0.5};
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);// bgr -> rgb
    cv::Mat img_float;
//    image.convertTo(img_float, CV_32F, 1.0 / 255);//归一化到[0,1]区间,
    cv::resize(image, img_float, cv::Size(img_W, img_H));

    std::vector<int64_t> dims = {1, img_H, img_W, 3};
    torch::Tensor img_var = torch::from_blob(img_float.data, dims, torch::kByte).to(device);//将图像转化成张量
    img_var = img_var.permute({0,3,1,2});//将张量的参数顺序转化为 torch输入的格式 NCHW
    img_var = img_var.toType(torch::kFloat);
    img_var = img_var.div(255);

    for (int i = 0; i < 3; i++) {
        img_var[0][i] = img_var[0][i].sub_(mean_[i]).div_(std_[i]);
    }
    img_var = torch::autograd::make_variable(img_var, false);
    return img_var;
}

int main(){

    std::string img_path = "../person.jpg";
    std::string model_path = "../model.pt";
    int img_W = 640;
    int img_H = 480;

    //加载模型到CPU
	clock_t t1 = clock();
    torch::DeviceType device_type;
    device_type = torch::kCPU;
    torch::Device device(device_type);
    torch::jit::script::Module module = torch::jit::load(model_path);//
    module.to(torch::kCPU);
    module.eval();
    std::cout<<"Model loading complete"<<std::endl;
	clock_t t2 = clock();
	std::cout << " loading model time: " << (t2 - t1) * 1.0 / CLOCKS_PER_SEC * 1000 << std::endl;
	int cnt = 0;
	int sum = 0;
	while(cnt < 10) {
		clock_t t2 = clock();
		std::string img_path = "../person" + std::to_string(cnt) + ".jpg";
		std::cout<<img_path<<std::endl;
		cv::Mat image = cv::imread(img_path, cv::ImreadModes::IMREAD_COLOR);
    	if (image.empty())
        	std::cout<<"Can not load image!"<<std::endl;
    	torch::Tensor img_var = process(image, device, img_W, img_H);
		clock_t t3 = clock();

	    //forward
	    torch::Tensor result = module.forward({img_var}).toTensor();

		clock_t t4 = clock();
	    result = result.squeeze();//删除1的维度
	    result = result.mul(255).to(torch::kU8) ;
	    cv::Mat pts_mat(cv::Size(img_W, img_H), CV_8U, result.data_ptr());
	    std::string dst = "../matting_person" + std::to_string(cnt) + ".jpg";
	    cv::imwrite(dst, pts_mat);
	    std::cout<<"matting complete!"<<std::endl;
		clock_t t5 = clock();

		std::cout << " image pre process time: " <<(t3 - t2) * 1.0 / CLOCKS_PER_SEC * 1000 << std::endl;
		std::cout << " inference time: " << (t4 - t3) * 1.0 / CLOCKS_PER_SEC * 1000 << std::endl;
		std::cout << " image pre process time: " <<(t5 - t4) * 1.0 / CLOCKS_PER_SEC * 1000 << std::endl;
		std::cout << " total time: " << (t5- t2) * 1.0 / CLOCKS_PER_SEC * 1000 << std::endl;

		if(cnt > 0) sum += (t5- t2) * 1.0 / CLOCKS_PER_SEC * 1000;
		cnt++;
	}
	std::cout<<sum / 9.0<<std::endl; //第一帧由于加载的原因速度很慢不算在内
	getchar();
	return 0;
}