下載libtorch的包
這裏下載的是cpu版本
Download here (Pre-cxx11 ABI):
https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-1.4.0%2Bcpu.zip
Download here (cxx11 ABI):
https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.4.0%2Bcpu.zip
也可以下載gpu版本(cuda10.1):
Download here (Pre-cxx11 ABI):
https://download.pytorch.org/libtorch/cu101/libtorch-shared-with-deps-1.4.0.zip
Download here (cxx11 ABI):
https://download.pytorch.org/libtorch/cu101/libtorch-cxx11-abi-shared-with-deps-1.4.0.zip
轉換模型
先使用torchvision自帶的resnet小試一下:
import torch
import torchvision
# An instance of your model.
model = torchvision.models.resnet18(pretrained=True)
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("traced_resnet18_model.pt")
In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
In[2]: output[0, :5]
Out[2]: tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)
在C ++中加載
- 使用C++編譯器(CLion)新建一個工程libtorch_test1
- 解壓libtorch-shared-with-deps-1.4.0+cpu.zip,把裏面的libtorch文件夾複製到工程目錄中
- 修改main.cpp 如下:
#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]) {
if (argc != 2) {
std::cerr << "usage: example-app <path-to-exported-script-module>\n";
return -1;
}
torch::jit::script::Module module;
try {
// Deserialize the ScriptModule from a file using torch::jit::load().
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
std::cout << "ok\n";
}
- 修改Cmakelists.txt如下:
cmake_minimum_required(VERSION 3.15)
project(libtorch_test1)
list(APPEND CMAKE_PREFIX_PATH "。/libtorch")
find_package(Torch REQUIRED)
set(CMAKE_CXX_STANDARD 14)
add_executable(libtorch_test1 main.cpp)
target_link_libraries(libtorch_test1 "${TORCH_LIBRARIES}")
set_property(TARGET libtorch_test1 PROPERTY CXX_STANDARD 14)
- 編譯
mkdir build
cd build
cmake ..
make
- 把轉換後的模型traced_resnet18_model.pt 複製過來,執行
./libtorch_test1 traced_resnet18_model.pt
將會打印出 OK 即爲成功。 - 在代碼中添加vector輸入,修改後的main.cpp 如下:
#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]) {
if (argc != 2) {
std::cerr << "usage: example-app <path-to-exported-script-module>\n";
return -1;
}
torch::jit::script::Module module;
try {
// Deserialize the ScriptModule from a file using torch::jit::load().
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
// 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';
std::cout << "ok\n";
}
- 重新編譯,執行
root@4b5a67132e81:/example-app/build# make
Scanning dependencies of target example-app
[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
[100%] Linking CXX executable example-app
[100%] Built target example-app
root@4b5a67132e81:/example-app/build# ./example-app traced_resnet_model.pt
-0.2698 -0.0381 0.4023 -0.3010 -0.0448
[ Variable[CPUFloatType]{1,5} ]
- 成功啦
參考: https://pytorch.org/tutorials/advanced/cpp_export.html#loading-a-pytorch-model-in-c