1.获取与安装
- 在主页面选择libtorch 下载
- 解压得到libtorch包(包括bin,include,lib,share,build-version)
- 安装就是把libtorch包放置到自己想要的地方(就放置在下载处也ok)
2.使用
- CMakelists.txt
project(test_pytorch)
set(CMAKE_CXX_STANDARD 11)
set(Torch_DIR /home/xxx/下载/libtorch/share/cmake/Torch) #指定libTorch位置(应该是有更好的办法安装)
#include_directories(${OpenCV_INCLUDE_DIRS} /home/xxx/下载/libtorch/include /home/xxx/下载/libtorch/include/torch/csrc/api/include)
find_package(OpenCV REQUIRED)
find_package(Torch REQUIRED) # 自动查找libTorch包
add_executable(test_pytorch test_pytorch.cpp)
target_link_libraries(test_pytorch ${OpenCV_LIBS} ${TORCH_LIBRARIES}) # 加入libTorch的库文件路径
- C++文件的编写
#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
using namespace std;
int main() {
//加载pytorch模型
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(
"/media/xxx/model.pt");
assert(module != nullptr);
//加载输入
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));
//指定执行位置
// Execute the model and turn its output into a tensor.
at::Tensor output = module->forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
cout << "Hello, world!" << endl;
}
3.选择计算Device(CPU或GPU)
最小代码
//设置Device
torch::DeviceType device_type; //设置Device类型
device_type = torch::kCUDA; //torch::kCUDA and torch::kCPU
torch::Device device(device_type, 0);
//把模型和数据都送到Device中去(数据和模型必须都在同一个device,结果也是)
module->to(device);
inputs.push_back(torch::ones({1, 3, 224, 224}).to(device));
完整代码:
#include <iostream>
#include "torch/script.h"
#include "torch/torch.h"
using namespace std;
int main() {
//加载pytorch模型
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(
"/media/xxxx/model.pt");//注意检查路径
assert(module != nullptr);
torch::DeviceType device_type; //设置Device类型
device_type = torch::kCUDA; //torch::kCUDA and torch::kCPU
torch::Device device(device_type, 0);
//模型转到GPU中去
module->to(device);
torch::cuda::is_available(); //判断是否支持GPU加速
// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}).to(device));
//指定执行位置
// Execute the model and turn its output into a tensor.
for (int i = 0; i < 100; i++) {
at::Tensor output = module->forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';
}
}
5.经验
- 学习新框架的最好方式就是查找源代码,最后一个Github的源代码帮助很大.
4.参考资料 - 官网说明 LOADING A PYTORCH MODEL IN C++
- 利用Pytorch的C++前端(libtorch)读取预训练权重并进行预测
- a github example
网友评论