前一篇博客讲了怎么用pytorch实现一个简单的图像分类器,这一篇的主要目的则是将前面训练好的模型用C++跑起来。
模型转换
第一步需要python环境下训练得到的pytorch模型转换为C++可读的模型,更具体的可以参考官方教程,这里用了最简单的通过Tracing的方法来进行模型转换:
import torch
from SimpleNet import SimpleNet
if __name__ == "__main__":
# An instance of your model.
best_model_path = 'epoch_462.pth'
model = SimpleNet()
model.load_state_dict(torch.load(best_model_path))
model.eval()
# An example input you would normally provide to the forward() method
example = torch.rand(1, 1, 90, 90)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
# Save model
traced_script_module.save("model.pt")
跑完以上代码应该会得到一个model.pt文件,这个文件就是接下来将用于C++下的模型文件。
C++程序代码
这里将整个预测过程写成了一个函数,输入为一个代表了一张灰度图片的二维数组,而在main函数里对此进行测试。代码如下:
#include #include #include // The function to predict the class of an image
int predict(void* img)
{
std::string model_path = "model.pt";
// load the model
std::shared_ptrmodule = torch::jit::load(model_path);
// Change the Image into Tensor for prediction
torch::Tensor tensor_image = torch::from_blob(img, { 1, 1, 90, 90 }, torch::kInt);
tensor_image = tensor_image.toType(torch::kFloat);
// distribution between 0 and 1
tensor_image = tensor_image.div(255);
// normalize, value between -1 and 1
tensor_image = tensor_image.sub(0.5);
tensor_image = tensor_image.div(0.5);
// predict
torch::Tensor result = module->forward({ tensor_image }).toTensor();
// get the class index
auto max_result = result.max(1, true);
auto max_index = std::get<1>(max_result).item();
return int(max_index);
}
int main()
{
// A test image whose class is bubble
//int img[90][90] = ...a 2-d array here...;
int pred = predict(img);
std::cout << "The predicted result is : " << pred << std::endl;
while (1);
return 0;
}
与python代码对应的,这里同样需要对输入图像做包括归一化操作在内的预处理,然后才能调用模型进行预测。上面的代码保存为一个example-app.cpp文件,保存在新建的一个SimpleNet文件夹下。
编译之前的准备
接着在SimpleNet文件夹下新建一个CMakeLists.txt文件并写入以下内容:
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(SimpleNet)
find_package(Torch REQUIRED)
add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 11)
Linux下用C++调用pytorch模型
首先从官网下载Linux版本的libtorch并解压:
Linux版本的libtorch
然后进入SimpleNet文件夹,依次执行以下命令:
mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=~/cmake/libtorch ..(这里libtorch路径需要根据实际解压路径作修改)
make
然后把model.pt文件复制到build文件夹里面去,便可以用以下命令执行了:
./example-app
Windows下用C++调用pytorch模型
首先从官网下载Windows版本的libtorch并解压:
Windows版本的libtorch
参考这篇博客,在SimpleNet文件夹中新建build文件夹,进入build文件夹,在命令行下输入:
cmake -DCMAKE_PREFIX_PATH=path\libtorch -DCMAKE_BUILD_TYPE=Release -G "Visual Studio 15 Win64" ..
其中DCMAKE_PREFIX_PATH修改为解压后的libtorch路径,Visual Studio 15代表的是VS2017所用的vc版本为15,如果用的VS2015则修改15为14,注意最后有..这两个点。
接着将model.pt复制到build文件夹中,进入build文件夹,用VS2017打开example-app.vcxproj文件,把example-app设置为启动项目,执行example-app.cpp,这时会报找不到xxx.dll的错,把libtorch/lib下的c10.dll、caffe2.dll、libiomp5md.dll和torch.dll复制到Release/Debug文件夹中去再重新执行即可(可能需要在debug模式下执行一次,再在release模式下执行一次)。