libtorch 调用模型方法

news2024/11/26 17:31:46

libtorch 是 pytorch 的 C++ 版本，支持 CPU 端和 GPU 端的模型部署。相较于用 tensorrt 部署 pytorch 模型，用 libtorch的优势在于：pytorch 和 libtorch 同属一个生态，API 语句比较接近，并且不会出现某网络层不支持的问题。这里咱们来看一下 libtorch 怎么调用模型进行推理的。

这里咱们以 resnet50 为例。

首先加载 resnet50 模型，并导出 trace 模型。

# 准备一个 py 脚本，里面内容如下
import torch
import torchvision

model = torchvision.models.resnet50(pretrained=False)
model = model.eval().cuda()
input_data = torch.randn(1, 3, 224, 224).cuda()

# export trace model
traced_script_model = torch.jit.trace(model, input_data)
output = traced_script_model(data)
traced_script_model.save('resnet50.pt')

print(output)

这里就会导出 trace 模型 resnet50.pt。然后编写推理工程。

先写 infer.cpp：

// infer.cpp
#include "torch/torch.h"
#include "torch/script.h"
#include <iostream>

int main(){
  // torch::Tensor tensor = torch::ones(3);
  // std::cout << tensor << std::endl;
  torch::jit::script::Module module;
  module = torch::jit::load("~/resnet50.pt");    // 导入前面生成的trace模型
  module.to(at::kCUDA);                          // 放到GPU上执行

  // 构建输入张量
  std::vector<torch::jit::IValue> inputs;
  inputs.push_back(torch::ones({1, 3, 224, 224}).to(at::kCUDA));

  // 执行推理
  at::Tensor output = module.forward(inputs).toTensor();

  std::cout << output << std::endl;
}

然后编写 CMakeList.txt：

cmake_minimum_required(VERSION 2.8 FATAL_ERROR)
find_package(PythonInterp REQUIRED)

project(infer_resnet50)

set(Torch_DIR ~/libtorch/share/cmake/Torch)  #你解压的libtorch的绝对路径
find_package(Torch REQUIRED)

set(CMAKE_CXX_FLAGS "${CAMKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
set(CUDA_INCLUDE_DIRS "/usr/local/cuda/include")

add_executable(infer_resnet50 infer.cpp)

#link libtorch .a .so
target_link_libraries(infer_resnet50 "${TORCH_LIBRARIES}")
target_include_directories(infer_resnet50 PRIVATE CUDA_INCLUDE_DIRS)

#
set_property(TARGET infer_resnet50 PROPERTY CXX_STANDARD 14)

开始编译 & 执行：

mkdir build

cd build

cmake ..

make -j8

正常执行输出结果：