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

前言
在我这里的实验结论表明，在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，肯定首先需要一个已经训练好模型，这里我使用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倍！效果还是很显著的，具体为什么会那么快，大概归为这几点：
1.ONNX-TensorRT将ONNX模型转化为TensorRT能够读懂的模型；
2.TensorRT将上一步转化后的模型进行了修改，融合了部分操作层和运算步骤，形成了新的融合后的模型
3.将融合后的模型进一步序列化到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的量化对精度的损失不是很高，是目前已经比较成熟的量化技术之一了。具体的量化步骤这里先不进行介绍了，TensorRT的INT8的实现可以查看Nvidia相关的PPT：http://ondemand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf
但是PPT只是讲了大概的量化流程，在TensorRT中量化是闭源的，目前只支持Caffe和TensorFlow模型的INT8量化，而ONNX模型的则暂未支持。
参考文献
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.csdn.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

转自: https://oldpan.me/archives/tensorrt-code-toturial-1