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

发表于由
阅读数 4,020

友情提示：

  • 阅读本文需要您已经掌握Pytorch的Python用法，并掌握C++语言。
  • 推荐使用Ubuntu/Mac系统实验（cmake可以自动找到已安装的opencv）。
  • 本实验需要已安装好opencv和pytorch 1.0，C++编译环境（Ubuntu需要g++，Mac需要XCode）和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测试模型

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 <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

int main(int argc, const char* argv[]) {
    if (argc != 2) {
      std::cerr << "usage: example-app <path-to-exported-script-module>\n";
      return -1;
    }
    // 加载导出的模型和权重
    std::shared_ptr<torch::jit::script::Module> 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<int64_t> 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<torch::jit::IValue> inputs;
    inputs.emplace_back(tensor_image);

    // 推理
    at::Tensor result = module->forward(inputs).toTensor();
    // 将输出张量转换为vector
    auto my_tensor = result;
    std::vector<float> vec(my_tensor.data<float>(), my_tensor.data<float>() + 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;
}

这段程序里比较难找的是将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的安装路径。

  • 执行编译成功的程序 
    ./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