Python1.0用C++实现残差网络图像分类

友情提示：

• 阅读本文需要您已经掌握Pytorch的Python用法，并掌握C++语言。
• 推荐使用Ubuntu/Mac系统实验（cmake可以自动找到已安装的opencv）。
• 本实验需要已安装好opencv和pytorch 1.0，C++编译环境（Ubuntu需要g++，Mac需要XCode，Windows需要Visual Studio）和cmake。

---

Pytorch 1.0已经于近日推出，其中一个亮点功能是支持将python训练的模型导出到C++进行推理。相比于目前流行的caffe训练模型+opencv dnn模块推理，pytorch从Python训练到C++部署提供了一体化的方案，可谓攻城狮的福音。

Python导出模型

根据Pytorch官网教程，我们先导出残差网络Resnet18的模型和预训练权重：

# coding=utf-8
import torch
import torchvision
from torchvision import transforms
from PIL import Image
import json
import cv2

# 初始化模型
model = torchvision.models.resnet18(pretrained=True)
model.eval() #将模型置为推理状态

# 随机生成一个输入张量
example = torch.rand(1, 3, 224, 224)

# 利用跟踪数据流的方法生成导出模型
traced_script_module = torch.jit.trace(model, example)
output = traced_script_module(torch.ones(1, 3, 224, 224))
print output.shape
print output[0, :5]
traced_script_module.save("model.pt")

这样在当前目录下会生成一个model.pt文件，包含了模型定义和权重。

Python测试模型

• 首先准备测试数据。在当前目录创建data目录，里面放入网上随便找的猫的图像，命名为cat.jpg。
• 下载1000类标签文件https://github.com/raghakot/keras-vis/blob/master/resources/imagenet_class_index.json，放在当前目录。
• 创建并执行以下Python代码：

import torch
import torchvision
from torchvision import transforms
import numpy as np
from PIL import Image
import json
import cv2

# An instance of your model.
model = torchvision.models.resnet18(pretrained=True)

######
## test with an image
with open('imagenet_class_index.json') as f:
    labels = json.load(f)
#print labels

img = cv2.imread('data/cat.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
#cv2.imshow('img', img)
#cv2.waitKey(0)
img = cv2.resize(img, (224, 224), img)

t = torch.from_numpy(img).float() / 255.0
t[:, :, 0] = (t[:, :, 0] - 0.485) / 0.229
t[:, :, 1] = (t[:, :, 1] - 0.456) / 0.224
t[:, :, 2] = (t[:, :, 2] - 0.406) / 0.225

t = t.unsqueeze(0)
t = t.permute(0, 3, 1, 2)
model.eval()
output = model.forward(t)
output = output[0].detach().numpy()
pred = output.argmax()
pred_label = labels[str(pred)]
print pred, pred_label

将输出：

285 [u'n02124075', u'Egyptian_cat']

表明输入图像的类别是285号“埃及猫”。

C++调用导出模型

• 创建以下test_resnet.cpp源码

#include  // One-stop header.
#include 
#include 
#include 

using namespace std;
using namespace cv;

int main(int argc, const char* argv[]) {
    if (argc != 2) {
      std::cerr << "usage: example-app \n";
      return -1;
    }
    try {
    	// 加载导出的模型和权重
    	std::shared_ptr module = torch::jit::load(argv[1]);

    	assert(module != nullptr);
    	std::cout << "Model loaded!\n";
    
    	// 利用OpenCV读取图像
    	Mat image_bgr, tmp_image, image;
    	image_bgr = imread("../data/cat.jpg");
    	// 将OpenCV默认的BGR通道顺序转换为RGB通道顺序
    	cvtColor(image_bgr, tmp_image, COLOR_BGR2RGB);
    	// 缩放为模型可以接受的 224x224
    	resize(tmp_image, tmp_image, cv::Size(224, 224), 0, 0, CV_INTER_LINEAR);
    	// 将像素值从[0, 255]缩放到 [0,1]区间
    	tmp_image.convertTo(image, CV_32F, 1.0 / 255, 0);
    	///////////// 关键：将图像转换为 Tensor 输入
    	// 定义张量，维度为 B x H x W x C
    	std::vector sizes = {1, image.rows, image.cols, 3};
    	at::TensorOptions options(at::kFloat);
    	// 从OpenCV的Mat将数据转换为 Tensor
    	at::Tensor tensor_image = torch::from_blob(image.data, at::IntList(sizes), options);
    	// 调整张量的通道顺序为 BxCXHXW
    	tensor_image = tensor_image.permute({0, 3, 1, 2});

    	// 将张量放入vector中准备输入 （因此模型可以接受多个输入）
    	std::vector inputs;
    	inputs.emplace_back(tensor_image);
    
    	// 推理
    	at::Tensor result = module->forward(inputs).toTensor();
    	// 将输出张量转换为vector
    	auto my_tensor = result;
    	std::vector vec(my_tensor.data(), my_tensor.data() + my_tensor.numel());
   		// 找到vector里数值最大的对应的下标，这就是分类的类别编号
    	int pred = -1;
    	double maxValue = -9999;
    	for(int i = 0; i < vec.size(); ++i) {
        	double v = vec[i];
       		if (v > maxValue) {
           		maxValue = v;
            	pred = i;
            	//cout << " i , maxValue = " << pred << ", " << maxValue << endl;
       	 	}
   	 	}
    	cout << "pred = " << pred << endl;
    }
    catch (exception & err) {
        cout << err.what() << endl;
    }
}

这段程序里比较难找的是将OpenCV的Mat转换为Tensor的那一段。官网教程里缺失了这样的例子。

• 创建CMakeList.txt准备编译：

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)
message(STATUS "OpenCV library status:")
message(STATUS "    version: ${OpenCV_VERSION}")
message(STATUS "    libraries: ${OpenCV_LIBS}")
message(STATUS "    include path: ${OpenCV_INCLUDE_DIRS}")


add_executable(test_resnet test_resnet.cpp)
target_link_libraries(test_resnet "${TORCH_LIBRARIES}" "${OpenCV_LIBS}")
set_property(TARGET test_resnet PROPERTY CXX_STANDARD 11)

• 利用cmake编译

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/Library/Python/2.7/site-packages/torch ..
make

注意 cmake后面的参数要替换为本机torch lib的安装路径。如果在Windows上，cmake命令应当替换成

cmake -DCMAKE_PREFIX_PATH=D:\Anaconda3\Lib\site-packages\torch;E:\libraries\opencv\build -DCMAKE_GENERATOR_PLATFORM=x64 ..

也就是在CMAKE_PREFIX_PATH中需要同时指定torch lib和opencv的路径。需要使用CUDA时需要指定生成目标平台为x64。

• 执行编译成功的程序

./test_resnet ../model.pt

将输出

Model loaded!
pred = 285

此分类结果与Python测试的结果一致。

参考

[1] https://pytorch.org/tutorials/advanced/cpp_export.html
[2] https://pytorch.org/cppdocs/api/classat_1_1_tensor.html#exhale-class-classat-1-1-tensor
[3] https://github.com/pytorch/pytorch/issues/14330