利用TensorRT实现神经网络提速(读取ONNX模型并运行)

关于本文章的最新更新请查看:oldpan博客

前言

这篇文章接着上一篇继续讲解如何具体使用TensorRT。

在之前已经写到过一篇去介绍什么是TensorRT：利用TensorRT对深度学习进行加速，这篇文章中大概已经基本讨论了TensorRT究竟是个什么东西以及怎么使用它。

而在这篇文章中我们主要介绍如何使用它在我们的实际任务中进行加速。

在我这里的实验结论表明，在FP32的精度下，使用TensorRT和不使用TensorRT在GPU上运行的速度比大概为3:1，也就是在我这个模型为前提条件下，TensorRT在GPU端使我的模型速度提升了3倍(不同模型不同显卡不同构架提升速度不同)。

TensorRT具备的功能

目前TensorRT的最新版本是5.0,TensorRT的发展其实已经有一段时间了,支持转化的模型也有caffe、tensorflow和ONNX了，我们要知道，TensorRT是有自己的模型框架的，我们首先先其他训练好的框架通过转化代码转化为TensorRT的代码才可以使用。TensorRT对Caffe模型的支持度最高，同时也支持将Caffe模型转化为int8精度。

而ONNX模型的转化则是近半年来的实现成果，目前支持了大部分的运算(经过测试，我们平常使用的90%的模型都可以使用ONNX-TensorRT来进行转化)。唯一遗憾的是ONNX模型目前还不支持int8类型的转化。

为什么需要转化，因为TensorRT只是一个可以在GPU上独立运行的一个库，并不能够进行完整的训练流程，所以我们一般是通过其他的神经网络框架(Pytorch、TensorFlow)训练然后导出模型再通过TensorRT的转化工具转化为TensorRT的格式再去运行。

TensorRT的优点在上一篇文中已经说过了，这里就不赘述。如果遇到关于TensorRT安装和基本概念请看上一篇文章：利用TensorRT对深度学习进行加速。

这一篇我们具体聊聊TensorRT的内在，以及我们该如何使用它。

利用TensorRT

我们在安装好TensorRT后(安装过程见上一篇文章)，对于我们来说，我们要使用TensorRT，肯定首先需要一个已经训练好模型，这里我使用ONNX，因为我自己经常使用的框架是Pytorch，所以我利用Pytorch导出了ONNX模型。

device = torch.device('cuda:0')

body = create_body(mobilenetv2(pretrained=False), -1)
nf = num_features_model(body) * 2   # Here we get the output channel from last layer
head = create_head(nf, 3, None, ps=0.5, bn_final=None)
model = nn.Sequential(body, head)

state = torch.load('new-mobilenetv2-128_S.pth', map_location=device)
model.load_state_dict(state['model'], strict=True)

example = torch.rand(1, 3, 128, 128).cuda()
model.to(device)

# 导出onnx模型
torch_out = torch.onnx.export(model,
                              example,
                              "new-mobilenetv2-128_S.onnx",
                              verbose=True,
                              export_params=True
                              )

上面的代码即展示了我的导出过程，利用改进后的mobilenetv2模型，然后读取.pth版的权重，最后再导出onnx，这一步骤具体解释官方都有的，如果不懂可以到官方教程中去查阅。

这样，我就导出了ONNX版本的模型：new-mobilenetv2-128_S.onnx

这里建议使用netron来可视化我们的模型：

如上所示是我刚才导出模型的可视化效果，我们可以看到模型图中的操作名称和我们一般使用的略有些区别，例如Clip就代表我们平时使用的Relu。从右侧的介绍栏中可以看到ONNX的版本是v3，op版本是v9，总共的操作数为92个。我们需要注意这些版本与支持的操作运算有着密切的关系。

准备显卡

上面我们已经导出了我们需要的ONNX模型，现在我们就要开始使用TensorRT了，但是需要注意，TensorRT只能用在GPU端，在纯CPU上是跑不了的，我们需要一张支持相关运算的显卡。在这里我是1080TI，1080TI支持fp32和int8精度的运算，而最新出的RTX2080TI系列则支持fp16，关于显卡计算能力和支持的运算可以看：新显卡出世，我们来谈谈与深度学习有关的显卡架构和相关技术。

显卡准备好，还有相关驱动也要安装好，具体步骤可以查看开头提到的那一篇文章。

TensorRT程序运行

首先我们修改一段官方的Sample(sampleOnnxMNIST)，大概步骤是使用ONNX-TensorRT转换工具将ONNX模型进行转换，然后使用TensorRT构建模型并运行起来。

省略掉代码中的其他的部分(想看完整的代码可以直接查看官方的例子)，这里只展示了修改后的main函数的部分内容：

IHostMemory* trtModelStream{nullptr};
// 这里读入刚才导出的模型
onnxToTRTModel("new-mobilenetv2-128_S.onnx", 1, trtModelStream);
assert(trtModelStream != nullptr);

// 利用Opencv设置输入信息，引入Opencv的头文件
cv::Mat src_host(cv::Size(128,128),CV_32FC3);

// deserialize the engine
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);

ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream->data(), trtModelStream->size(), nullptr);
assert(engine != nullptr);
trtModelStream->destroy();
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);

float prob[OUTPUT_SIZE];
// 将输入图像数据的数据格式由0-255转化为0-1
for (int i = 0; i < INPUT_H * INPUT_W * 3; ++i)
    data[i] = float(src_host.data[i] / 255.0);

// 这里我测试了一下时间
auto startTime = std::chrono::high_resolution_clock::now();

for(int i = 0 ; i < 10000 ; i++)
    doInference(*context, data, prob, 1);

auto endTime = std::chrono::high_resolution_clock::now();
float totalTime = std::chrono::duration<float, std::milli>(endTime - startTime).count();

std::cout << "Time used one image (measured by chrono):" << totalTime/10000 << " ms" << std::endl;

// destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();

上面这段代码打算测试了利用TensorRT去跑ONNX模型的速度。

需要注意一点，在测试GPU所运行的时候我们需要用到下面的函数使GPU和CPU保持同步，这样我们测GPU的运行时间才会精准，当然在TensorRT的例程中已经利用下面这个语句进行了同步操作。

cudaStreamSynchronize()：这个方法接受一个stream ID，它将阻止CPU执行直到GPU端完成相应stream ID的所有CUDA任务，但其它stream中的CUDA任务可能执行完也可能没有执行完。

接下来我们开始编译，由于官方提供的示例程序中使用的是makefile文件，不利于我们之后的修改，所以为了方便我们根据官方提供的makefile文件编写成了CmakeList版本，方便以后修改：

cmake_minimum_required(VERSION 3.12)
project(tensorrt)

#set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")	# -std=gnu++11

set(CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER})

# 查找CUDA
find_package(CUDA)

# 在这里修改我们显卡的计算能力 这里我是sm_61
set(
        CUDA_NVCC_FLAGS
        ${CUDA_NVCC_FLAGS};
        -O3
        -gencode arch=compute_61,code=sm_61
)

set(PROJECT_OUTPUT_DIR  ${PROJECT_BINARY_DIR}/${CMAKE_SYSTEM_PROCESSOR})
set(PROJECT_INCLUDE_DIR ${PROJECT_OUTPUT_DIR}/include)

file(MAKE_DIRECTORY ${PROJECT_INCLUDE_DIR})
file(MAKE_DIRECTORY ${PROJECT_OUTPUT_DIR}/bin)

set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_OUTPUT_DIR}/bin)  # .exe .dll
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_OUTPUT_DIR}/lib)  # .dll .so
set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${PROJECT_OUTPUT_DIR}/lib)  # .lib .a

include_directories(${PROJECT_INCLUDE_DIR})
include_directories(${PROJECT_SOURCE_DIR}/include)

file(GLOB Sources *.cpp)
file(GLOB Includes include/*.h)

foreach(include ${Includes})
    message("-- Copying ${include}")
    configure_file(${include} ${PROJECT_INCLUDE_DIR} COPYONLY)
endforeach()

find_package(Protobuf)

if(PROTOBUF_FOUND)
    message(STATUS "    version: ${Protobuf_VERSION}")
    message(STATUS "    libraries: ${PROTOBUF_LIBRARIES}")
    message(STATUS "    include path: ${PROTOBUF_INCLUDE_DIR}")
else()
    message(WARNING "Protobuf not found, onnx model convert tool won't be built")
endif()

set(TENSORRT_ROOT /home/prototype/Downloads/TensorRT-5.0.2.6)
find_path(TENSORRT_INCLUDE_DIR NvInfer.h
        HINTS ${TENSORRT_ROOT} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES include)
MESSAGE(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
find_library(TENSORRT_LIBRARY_INFER nvinfer
        HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
find_library(TENSORRT_LIBRARY_INFER_PLUGIN nvinfer_plugin
        HINTS  ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_INFER_PLUGIN})
MESSAGE(STATUS "Find TensorRT libs at ${TENSORRT_LIBRARY}")
find_package_handle_standard_args(
        TENSORRT DEFAULT_MSG TENSORRT_INCLUDE_DIR TENSORRT_LIBRARY)
if(NOT TENSORRT_FOUND)
    message(ERROR
            "Cannot find TensorRT library.")
endif()

LINK_LIBRARIES("/home/prototype/Downloads/TensorRT-5.0.2.6/lib/libnvonnxparser.so")

find_package(OpenCV REQUIRED)

cuda_add_executable(tensorrt benchmark.cpp)

target_include_directories(tensorrt PUBLIC ${CUDA_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR})
target_link_libraries(tensorrt ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY} ${CUDA_CUBLAS_LIBRARIES} ${CUDA_cudart_static_LIBRARY} ${OpenCV_LIBS})

cmake文件主要注意的几点是cuda和TensorRT动态链接库的查找，只要这几个必备的动态链接库使用cmake查找到就可以编译成功，另外由于我是用了Opencv，所以也将Opencv也加入了编译。

编译后运行，发现利用TensorRT在FP32精度下跑相同模型比在Pytorch的C++端跑几乎快了3倍！效果还是很显著的，具体为什么会那么快，大概归为这几点：

• ONNX-TensorRT将ONNX模型转化为TensorRT能够读懂的模型；
• TensorRT将上一步转化后的模型进行了修改，融合了部分操作层和运算步骤，形成了新的融合后的模型
• 将融合后的模型进一步序列化到GPU中，并对特定GPU进行了优化操作，使推断更快速。

上面的步骤其实可以通过官方的开发手册中清晰地看到：

TensorRT是闭源的，官方只是提供了如何去使用它但是并没有提供我们源代码，上述提升速度的核心要点是模型融合和操作简化(因为我的1080Ti不支持FP16所以默认使用的FP32，而且ONNX-TensorRT目前不支持int8的转换)，其中支持的模型融合可以查看官方提供的列表，最常用的即conv+bn+relu也就是下图中红箭头指向的地方：

其实int8提速的效果也是很大的，在官方有一个简单的MNIST的INT8的例子，在MNIST这个简单的任务中提速可以达到20%(任务越大提升速度越明显)，在其他的任务上应该也有不错的提升效果。

TensorRT的已知BUG

TensorRT虽然一直在更新，但是其BUG还是不少的，下图中遇到的问题，我在TensorRT-5.0中也遇到了：

也就是我们在使用TensorRT在进行模型转化的过程中会出现：Signal: SIGSEGV (Segmentation fault)。

具体的错误代码发生在以下几个语句中，错误原因很有可能是已经释放的内存再次被释放。

engine->destroy();
network->destroy();
builder->destroy();

TensorRT知识简介

TensorRT的资料除了官方资料之外并没有其余的资料了，所幸TensorRT的代码库中注释很详细，我们可以通过TensorRT的头文件代码尽可能地了解TensorRT这个库。

TensorRT中的数据类型

目前TensorRT支持的数据类型有以下四种，除了我们平常使用的kFLOAT(单精度浮点型)，TensorRT还支持半精度浮点型、量化INT8类型以及INT32类型：

这四种类型用于权重信息和张量信息中，例如我们在使用TensorRT库设计网络层的时候就需要注明：

// ITensor是TensorRT中定义张量的类
// 下面我们直接用TensorRT中的方法定义了该网络的输入 其中dt为参数 由 DataType dt 传递进来
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});

利用 INetworkDefinition 定义神经网络

在sampleMNISTAPI这个官方例程中，

// Creat the engine using only the API and not any parser.
ICudaEngine* createMNISTEngine(unsigned int maxBatchSize, IBuilder* builder, DataType dt)
{
    INetworkDefinition* network = builder->createNetwork();

    // Create input tensor of shape { 1, 1, 28, 28 } with name INPUT_BLOB_NAME
    ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});
    assert(data);

    // Create scale layer with default power/shift and specified scale parameter.
    const float scaleParam = 0.0125f;
    const Weights power{DataType::kFLOAT, nullptr, 0};
    const Weights shift{DataType::kFLOAT, nullptr, 0};
    const Weights scale{DataType::kFLOAT, &scaleParam, 1};
    IScaleLayer* scale_1 = network->addScale(*data, ScaleMode::kUNIFORM, shift, scale, power);
    assert(scale_1);

    // Add convolution layer with 20 outputs and a 5x5 filter.
    std::map<std::string, Weights> weightMap = loadWeights(locateFile("mnistapi.wts"));
    IConvolutionLayer* conv1 = network->addConvolution(*scale_1->getOutput(0), 20, DimsHW{5, 5}, weightMap["conv1filter"], weightMap["conv1bias"]);
    assert(conv1);
    conv1->setStride(DimsHW{1, 1});

    // Add max pooling layer with stride of 2x2 and kernel size of 2x2.
    IPoolingLayer* pool1 = network->addPooling(*conv1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
    assert(pool1);
    pool1->setStride(DimsHW{2, 2});

    // Add second convolution layer with 50 outputs and a 5x5 filter.
    IConvolutionLayer* conv2 = network->addConvolution(*pool1->getOutput(0), 50, DimsHW{5, 5}, weightMap["conv2filter"], weightMap["conv2bias"]);
    assert(conv2);
    conv2->setStride(DimsHW{1, 1});

    // Add second max pooling layer with stride of 2x2 and kernel size of 2x3>
    IPoolingLayer* pool2 = network->addPooling(*conv2->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
    assert(pool2);
    pool2->setStride(DimsHW{2, 2});

    // Add fully connected layer with 500 outputs.
    IFullyConnectedLayer* ip1 = network->addFullyConnected(*pool2->getOutput(0), 500, weightMap["ip1filter"], weightMap["ip1bias"]);
    assert(ip1);

    // Add activation layer using the ReLU algorithm.
    IActivationLayer* relu1 = network->addActivation(*ip1->getOutput(0), ActivationType::kRELU);
    assert(relu1);

    // Add second fully connected layer with 20 outputs.
    IFullyConnectedLayer* ip2 = network->addFullyConnected(*relu1->getOutput(0), OUTPUT_SIZE, weightMap["ip2filter"], weightMap["ip2bias"]);
    assert(ip2);

    // Add softmax layer to determine the probability.
    ISoftMaxLayer* prob = network->addSoftMax(*ip2->getOutput(0));
    assert(prob);
    prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);
    network->markOutput(*prob->getOutput(0));

    // Build engine
    builder->setMaxBatchSize(maxBatchSize);
    builder->setMaxWorkspaceSize(1 << 20);
    samplesCommon::enableDLA(builder, gUseDLACore);
    ICudaEngine* engine = builder->buildCudaEngine(*network);

    // Don't need the network any more
    network->destroy();

    // Release host memory
    for (auto& mem : weightMap)
    {
        free((void*) (mem.second.values));
    }

    return engine;
}

直接读取TensorRT原生支持的网络

安装支持CUDA的OpenCV

我们平时安装的OpenCV大多数只是在CPU环境下运行的，直接编译的话并没有CUDA的支持。强行使用CUDA模块只会报错。另外，从OpenCV-4.0之后，CUDA模块被移动到了opencv_contrib中，默认的源码包是不带CUDA的:

因此如果我们只下载OpenCV源码编译是不行的，必须加上contrib模块，下载好之后，在Cmake编译时添加contrib模块的路径：-DOPENCV_EXTRA_MODULES_PATH=<opencv_contrib>/modules并开启-DWITH_CUDA=ON。

这样我们开启CUDA编译好之后就可以使用OpenCV的功能了，也就可以向TensorRT直接传递GPU图像数据了。

如果我们没有安装CUDA版本的OpenCV，我们可以使用TensorRT中定义CPU版本的mat图像格式然后转化为TensorRT可以接受的数据格式。

// 定义一个变量接受Mmat数据
float data[INPUT_H * INPUT_W * 3];

// 将OpenCV的图像数据转化为纯float数组
void Mat_to_CHW(float *data, cv::Mat &frame)
{
    assert(data && !frame.empty());
    unsigned int volChl = INPUT_H * INPUT_W;
//    unsigned int volImg = INPUT_H * INPUT_W * INPUT_C;

    for(int c = 0; c < INPUT_C; ++c)
    {
        for (unsigned j = 0; j < volChl; ++j)
            data[c*volChl + j] = float(frame.data[j * INPUT_C + c]) / 255.0;
    }

    return;
}

...
// 在这里将之前的data 传入GPU中进行运算
// DMA the input to the GPU,  execute the batch asynchronously, and DMA it back:
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * INPUT_H * INPUT_W * INPUT_C * sizeof(float), cudaMemcpyHostToDevice, stream));
...

TensorRT的量化INT8

INT8相比如FLOAT16来说，可以更好地优化内存的使用量，并且速度相比FLOAT16提升更大(拥有更低的延迟和更高的吞吐量)。但是前提我们的显卡需要有足够多的INT8运算单元，官方建议使用6.1和7.x计算能力GPU显卡，我们经常使用的1080TI就满足要求，其计算能力为6.1。拥有足够数量的INT8运算单元。

不懂计算能力的可以看这篇文章：新显卡出世，我们来谈谈与深度学习有关的显卡架构和相关技术。

Int8的精度范围：

	取值范围	最小正值
F32	-3.4 x 10^38 ~ +3.4 x 10 ^38	1.4 x 10^45
F16	-65504 ~ 65504	5.96 x 10^8
INT8	-127 ~ 128	1

INT8的量化对精度的损失不是很高，是目前已经比较成熟的量化技术之一了。

具体的量化步骤这里先不进行介绍了，TensorRT的INT8的实现可以查看Nvidia相关的PPT：http://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf

但是PPT只是讲了大概的量化流程，在TensorRT中量化是闭源的，目前只支持Caffe和TensorFlow模型的INT8量化，而ONNX模型的则暂未支持。

其他类似的落地技术

技术落地也是深度学习中比较重要的一环，目前已经存在了很多的落地技术：Glow、TVM、Tensor Comprehensions、ncnn等都是为了能够将我们已经实现的模型进行优化到移动设备和嵌入式设备当中，相信未来的两三年中深度学习的落地方面将会大大发展起来。

参考文献

https://mxnet.incubator.apache.org/tutorials/tensorrt/inference_with_trt.html
https://devtalk.nvidia.com/default/topic/1030042/jetson-tx1/loading-of-the-tensorrt-engine-in-c-api/
https://petewarden.com/2015/05/23/why-are-eight-bits-enough-for-deep-neural-networks/
https://blog.****.net/zhangjunhit/article/details/84562334
https://www.jianshu.com/p/43318a3dc715
https://arleyzhang.github.io/articles/923e2c40/
https://mp.weixin.qq.com/s/F_VvLTWfg-COZKrQAtOSwg
https://mp.weixin.qq.com/s/wyqxUlXxgA9Eaxf0AlAVzg
https://elinux.org/Jetson/Performance
https://www.leiphone.com/news/201610/s2fwkopa5E1oCJxD.html
https://devtalk.nvidia.com/default/topic/1030567/tensorrt/tensorrt3-0-install-error-on-ubuntu-16-04-depends-cuda-cublas-9-0-but-it-is-not-installable-/
https://yq.aliyun.com/articles/600425?spm=a2c4e.11153940.blogcont497080.17.a3ac2f68x2E2bh
https://devtalk.nvidia.com/default/topic/1038826/tensorrt/layer-information-after-optimization-/post/5279684/#5279684