大家好，我是极智视界，本文介绍一下 libtorch 调用模型推理方法。
之前写了一篇《ubuntu 安装 libtorch》，所以关于 ubuntu 上安装 libtorch 的方法，有兴趣的同学可以自行查看。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
正常执行输出结果：
好了，以上分享了 libtorch 调用模型推理方法。希望我的分享能对你的学习有一点帮助。
