C++ 调用 Torch pt文件

Posted on | In , , |

[TOC]

install

https://pytorch.org/get-started/previous-versions/

https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model

python库 https://download.pytorch.org/whl/cu101/torch_stable.html

example-app.cpp

1
2
3
4
5
6
7
#include <torch/torch.h>
#include <iostream>

int main() {
  torch::Tensor tensor = torch::rand({2, 3});
  std::cout << tensor << endl;
}

Cmakelists.xml

1
2
3
4
5
6
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example-app)
find_package(Torch REQUIRED)
add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 11)

1| pth-2-pt

官方jit文档

​ https://pytorch.org/docs/1.5.1/jit.html

​ https://pytorch.org/docs/stable/generated/torch.jit.trace_module.html

  • 模型里面要严格遵守函数在init初始化，forward来使用

    看下面这个例子：

  • 对比torch.jit.trace功能，此处foward函数里面有依赖于输入数据的控制流，故必须使用torch.jit.script功能。

1
2
3
torch.jit.script(obj)
torch.jit.trace(obj, (args...))
torch.jit.trace_module()

torch.jit.trace:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import torch
import torchvision.models as models
import torch.nn as nn

# 载⼊预训练的型
model = models.squeezenet1_1(pretrained=False)
model.classifier[1] = nn.Conv2d(in_channels=512, out_channels=2, kernel_size=(1, 1), stride=(1, 1))
model.num_classes = 2
model.load_state_dict(torch.load('./model/model_squeezenet_utk_face_20.pth', map_location='cpu'))
model.eval()
print(model)
# model = models.resnet18(pretrained=True)
# print(model)

example = torch.rand(1, 3, 224, 224)
traced_script_module = torch.jit.trace(model, example)

output = traced_script_module(torch.ones(1, 3, 224, 224))
print(output)
# ----------------------------------
traced_script_module.save("./model/model_squeezenet_utkface.pt")

2| c++ load pt

由于涉及到图像的加载与处理，本⼈使⽤opencv进⾏读取和处理。
Tips:
训练过程中，采⽤PIL.Iamge加载图像(3通道 RGB)，然后Resize到224 x 224⼤⼩， 之后再进⾏ToTensor。因此使⽤C++ libTorch时候也需要按照上述过程对图像进⾏预处理。

  1. cv::imread()默认读取为三通道BGR，需要进⾏B/R通道交换，这⾥采⽤cv::cvtColor() 实现。
  2. 缩放cv::resize() 实现。
  3. opencv读取的图像矩阵存储形式：H x W x C, 但是pytorch中 Tensor的存储为：N x C x H x W, 因此需要进⾏变换，就是np.transpose()操作，这⾥使⽤tensor.permut({2,0,1})实现，效果是⼀样的。
  4. 数据归⼀化，采⽤tensor.div(255)实现。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
#include <torch/script.h> // One-stop header.
#include <opencv2/opencv.hpp>
#include <iostream>
#include <memory>
//https://pytorch.org/tutorials/advanced/cpp_export.html
string image_path = "/home/weipenghui/Project-dev/Cpp/testLibTorch2/image";

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/model_squeezenet_utkface.pt");
    assert(module != nullptr);
    std::cout << "ok\n";

    //输⼊图像
    auto image = cv::imread(image_path +"/"+ "7.jpg",cv::ImreadModes::IMREAD_COLOR);
    cv::Mat image_transfomed;
    cv::resize(image, image_transfomed, cv::Size(224, 224));
    cv::cvtColor(image_transfomed, image_transfomed, cv::COLOR_BGR2RGB);

    // 图像转换为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);

    // ⽹络前向计算
    // Execute the model and turn its output into a tensor.
    at::Tensor output = module->forward({tensor_image}).toTensor();
    auto max_result = output.max(1,true);
    auto max_index = std::get<1>(max_result).item<float>();
    if (max_index == 0){
        cv::putText(image, "male", cv::Point(50, 50), 1, 1,cv::Scalar(0, 255, 255));
    }else{
        cv::putText(image, "female", cv::Point(50, 50), 1, 1,cv::Scalar(0, 255, 255));
    }
    cv::imwrite("./result7.jpg", image);

    //cv::imshow("image", image);
    //cv::waitKey(0);
    //    at::Tensor prob = torch::softmax(output,1);
    //    auto prediction = torch::argmax(output, 1);
    //
    //    auto aa = prediction.slice(/*dim=*/0, /*start=*/0, /*end=*/2).item();
    //
    //    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/2) << '\n';
    //    std::cout << prob.slice(/*dim=*/1, /*start=*/0, /*end=*/2) << '\n';
    //

3| CMakeList

1
2
3
4
5
6
7
8
9
10
11
12
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)
#include_directories(./opencv-3.4.3/build/install_cv/include)
# set(CMAKE_PREFIX_PATH "./opencv-3.4.3/build/install_cv")
find_package(OpenCV REQUIRED)
set(CMAKE_PREFIX_PATH
        ./libtorch
        ./opencv-3.4.3/build/install_cv)
find_package(Torch REQUIRED)
add_executable(example-app main.cpp)
target_link_libraries(example-app ${TORCH_LIBRARIES} ${OpenCV_LIBS})
set_property(TARGET example-app PROPERTY CXX_STANDARD 11)

参考资料：

1、环境配置Pytorch Offical libtorch make demo

2、 模型转换 model To torchScript

3、模型调用LOADING A TORCHSCRIPT MODEL IN C++

C++ 调用libtorch，推理模型，very全面

VS2017 C++调用torch模型——libtorch

libtroch部署之torch.jit.script Module踩坑之旅

Demos

行人重识别PCB

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
#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;
}