Linux&Windows下用C++跑pytorch模型

 前一篇博客讲了怎么用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 <torch/script.h>

#include <iostream>
#include <memory>

// The function to predict the class of an image
int predict(void* img)
{
	std::string model_path = "model.pt";
	// load the model
	std::shared_ptr<torch::jit::script::Module> module = 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<float>();

	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模式下执行一次）。