歡迎關注我的公衆号 [極智視界],擷取我的更多筆記分享
大家好,我是極智視界,本文介紹一下 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
正常執行輸出結果:
好了,以上分享了 libtorch 調用模型推理方法。希望我的分享能對你的學習有一點幫助。