深度学习模型部署(1)-Pytorch转TensorRT(C++版)

xcyuyuyu

佛系生活，认真更新~

1. 前言

深度学习的应用落地流程包括了神经网络的训练与部署，我们熟悉的Tensorflow、Caffe、Pytorch等框架属于训练框架，用于神经网络的训练，在完成网络训练之后，我们需要对网络进行部署，这时候往往需要将网络迁移到部署框架上。

TensorRT是Nvidia开发的一个神经网络前向推理加速的C++库，提供了包括神经网络模型计算图优化、INT8量化、FP16低精度运算等神经网络前向推理优化的方法（更加具体的介绍可以参考官网： https:// developer.nvidia.com/te nsorrt ，初学者亦可参考深蓝学院的课程：《 CUDA入门与深度神经网络加速》）。目前TensorRT提供了C++与Python的API接口，本文中主要使用C++接口为例说明TensorRT框架的一般使用流程。

本文采用的实验流程为：Pytorch -> Onnx -> TensorRT。即首先将Pytorch模型转换为Onnx模型，然后通过TensorRT解析Onnx模型，创建TensorRT引擎及进行前向推理。

2. 准备工作

2.1 实验环境

GPU：1080ti；
Cuda：10.0；
Cudnn：7.6；
Pytorch：1.5；
TensorRT：7.0

2.2 TensorRT安装

在官网下载对应版本的TensorRT，我这里选择源码安装，所以我下载的是：

2. 安装TensorRT

按照官网安装指引进行安装即可

2.3 ONNX环境安装

>> sudo apt-get install libprotobuf-dev protobuf-compiler # protobuf is a prerequisite library
>> git clone --recursive https://github.com/onnx/onnx.git # Pull the ONNX repository from GitHub 
>> cd onnx
>> mkdir build && cd build 
>> cmake .. # Compile and install ONNX
>> make # Use the ‘-j’ option for parallel jobs, for example, ‘make -j $(nproc)’ 
>> make install

3. 示例详解

3.1 Python模型转Onnx

3.2 输入数据预处理idia官方提供的一个基于unet模型的大脑MRI图像分割的例子：

>> wget https://developer.download.nvidia.com/devblogs/speeding-up-unet.7z // Get the ONNX model and test the data
>> tar xvf speeding-up-unet.7z # Unpack the model data into the unet folder
>> cd unet
>> python create_network.py #Inside the unet folder, it creates the unet.onnx file

通过 create_network.py 脚本我们可以将pytorch模型转换为onnx模型，脚本的具体代码如下：

import torch
from torch.autograd import Variable
import torch.onnx as torch_onnx
import onnx
def main():
    input_shape = (3, 256, 256)
    model_onnx_path = "unet.onnx"
    dummy_input = Variable(torch.randn(1, *input_shape))
    model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
      in_channels=3, out_channels=1, init_features=32, pretrained=True)
    model.train(False)
    inputs = ['input.1']
    outputs = ['186']
    dynamic_axes = {'input.1': {0: 'batch'}, '186':{0:'batch'}}
    out = torch.onnx.export(model, dummy_input, model_onnx_path, input_names=inputs, output_names=outputs, dynamic_axes=dynamic_axes)
if __name__=='__main__':
    main()

pytorch1.5中提供了torch模型转onnx模型的函数，看代码可以很容易的理解pytorch转onnx的流程，更进一步的，可以在 torch.onnx 中了解到函数更加详细的介绍。

另外，如果大家想可视化onnx模型，这里推荐 Netron 这个工具，直接将onnx文件打开即可看到模型结构，十分方便。

3.3 输入数据预处理

在 Kaggle 中下载所有的图片，或者 https:// pan.baidu.com/s/137BDGP YckWRLkJVURzECkA ，提取码：myg4 。

任意选择所下载图片中的3张图片（图片名中没有_mask的），然后使用 prepareData.py 将图片进行预处理以及序列化并保存：

>> pip install medpy #dependency for utils.py file
>> mkdir test_data_set_0
>> mkdir test_data_set_1
>> mkdir test_data_set_2
>> python prepareData.py --input_image your_image1 --input_tensor test_data_set_0/input_0.pb --output_tensor test_data_set_0/output_0.pb   # This creates input_0.pb and output_0.pb
>> python prepareData.py --input_image your_image2 --input_tensor test_data_set_1/input_0.pb --output_tensor test_data_set_1/output_0.pb   # This creates input_0.pb and output_0.pb
>> python prepareData.py --input_image your_image3 --input_tensor test_data_set_2/input_0.pb --output_tensor test_data_set_2/output_0.pb   # This creates input_0.pb and output_0.pb

这样，我们完成了pytorch模型转onnx模型，以及输入数据预处理的步骤，下一步就是用TensorRT接口调用模型、创建引擎、并进行推理。

3.3 示例代码下载

首先下载官方提供的上述unet模型对应的TensorRT示例，并进行编译：

>> git clone https://github.com/parallel-forall/code-samples.git
>> cd code-samples/posts/TensorRT-introduction
>> make clean && make # Compile the TensorRT C++ code

在编译示例代码的时候，可能会遇到 NvInfer.h 文件或者其他TensorRT文件找不到的情况，此时可以在 code-samples/posts/TensorRT-introduction/Makefile 中加入TensorRT的路径：

将：
CXXFLAGS=-std=c++11 -DONNX_ML=1 -Wall -I$(CUDA_INSTALL_DIR)/include
LDFLAGS=-L$(CUDA_INSTALL_DIR)/lib64 -L$(CUDA_INSTALL_DIR)/lib64/stubs -L/usr/local/lib
CXXFLAGS=-std=c++11 -DONNX_ML=1 -Wall -I$(CUDA_INSTALL_DIR)/include -I/YoutPath/TensorRT-7.0.0.11/include
LDFLAGS=-L$(CUDA_INSTALL_DIR)/lib64 -L$(CUDA_INSTALL_DIR)/lib64/stubs -L/usr/local/lib -L/YourPath/TensorRT-7.0.0.11/targets/x86_64-linux-gnu/lib
(其中YourPath改为你自己的路径)

然后直接运行代码即可：

>> cd to code-samples/posts/TensorRT-introduction
>> ./simpleOnnx_1 path/to/unet/unet.onnx path/to/unet/test_data_set_0/input_0.pb # The sample application expects output reference values in path/to/unet/test_data_set_0/output_0.pb

运行成功之后，会在源目录下生成output.pgm文件，即为模型的输出。

3.4 主要代码分析

创建 CUDA Engine；

// Declare the CUDA engine
unique_ptr<ICudaEngine, Destroy<ICudaEngine>> engine{nullptr};
// Create the CUDA engine
engine.reset(createCudaEngine(onnxModelPath));

其中，创建 Engine 的函数为 createCudaEngine，具体代码如下：

nvinfer1::ICudaEngine* createCudaEngine(string const& onnxModelPath, int batchSize){
    unique_ptr<nvinfer1::IBuilder, Destroy<nvinfer1::IBuilder>> builder{nvinfer1::createInferBuilder(gLogger)};
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    unique_ptr<nvinfer1::INetworkDefinition, Destroy<nvinfer1::INetworkDefinition>> network{builder->createNetworkV2(explicitBatch)};
    unique_ptr<nvonnxparser::IParser, Destroy<nvonnxparser::IParser>> parser{nvonnxparser::createParser(*network, gLogger)};
    unique_ptr<nvinfer1::IBuilderConfig,Destroy<nvinfer1::IBuilderConfig>> config{builder->createBuilderConfig()};
    if (!parser->parseFromFile(onnxModelPath.c_str(), static_cast<int>(ILogger::Severity::kINFO)))
            cout << "ERROR: could not parse input engine." << endl;
            return nullptr;
    builder->setMaxBatchSize(batchSize);
    config->setMaxWorkspaceSize((1 << 30));
    auto profile = builder->createOptimizationProfile();
    profile->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kMIN, Dims4{1, 3, 256 , 256});
    profile->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kOPT, Dims4{1, 3, 256 , 256});
    profile->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kMAX, Dims4{32, 3, 256 , 256});    
    config->addOptimizationProfile(profile);
    return builder->buildEngineWithConfig(*network, *config);
}

2. 创建 execution context；

context 的功能是存储模型推理过程中的中间层激活值，创建的代码如下：

// Declare the execution context
unique_ptr<IExecutionContext, Destroy<IExecutionContext>> context{nullptr};
// Create the execution context
context.reset(engine->createExecutionContext());

3. 执行前向推理；

前向推理的代码如下所示，推理的过程为：1）将存储在CPU端的输入数据copy到GPU上；2）context->enqueueV2(bindings, stream, nullptr) 执行GPU端推理；3）将GPU端的输出copy回CPU端。

void launchInference(IExecutionContext* context, cudaStream_t stream, vector<float> const& inputTensor, vector<float>& outputTensor, void** bindings, int batchSize)
    int inputId = getBindingInputIndex(context);