TensorRT

1. TensorRT 简介

TensorRT 是一个前向推理框架。在推理过程中，基于TensorRT 的应用程序的执行速度可以比 CPU 平台速度快 40 倍。

  • 不同的硬件需要匹配不同的 cuda库，然后还需要进行测试， 比如选核等操作
  • TensorRT 以 NVIDIA 的并行编程模型 CUDA 为基础构建而成。
  • TensorRT 针对多种深度学习推理应用的生产部署提供 INT8 和 FP 16 优化

TensorRT 框架支持大多数的 DeepLearning 框架

TensorRT 核心优化方法

2. TensorRT OSS

(1) TensorRT OSS

TensorRT 分为开源和闭源两部分

开源部分：https://github.com/NVIDIA/TensorRT

​ NVIDIA TensorRT 的开源软件(OSS) 组件：其中包括 TensorRT 插件和解析器(caffe 和 ONNX) 的资源， 以及演示 TensorRT 平台用法和功能的示例应用程序。 这些开源软件组件是 TensorRT General Availibility(GA) 发行版的子集， 具有一些扩展和错误修复。

闭源版本： 量化、推理计算、kernel 查找等

(2) 安装流程

git 获取完成的 OSS 项目包

git clone -b https://github/nvidia/TensorRT TensorRT
cd TensorRT
git submodule update --init --recursive # 安装一些子模块: 比如 cub、protobuf、onnx、pybind11 等
export TRT_SOURCE='pwd'

安装 GA 包

cd ~/Downloads
tar -xvzf TensorRT-7.2.1.6.CentOS-7.6.x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz
export TRT_RELEASE=`pwd`/TensorRT-7.2.1.6

OSS 项目编译

cd $TRT_SOURCE
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_RELEASE/lib -DTRT_OUT_DIR=`pwd`/out # 会将 OSS 版本覆盖原有版本
# 输出二进制库 lib 和 可执行文件(示例代码)
make -j$(nproc)
(3) TensorRT OSS 目录
  • parser
  • plugin
  • examples

3. TensorRT 开发

（1）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)
（2）计算图优化：onnx 优化

常见的处理：

  • reshape 的时候的很多冗余算子的问题

  • 输入的归一化：去掉，放在代码逻辑里面

  • conv_bn 的合并
  • 一些无法 boardcast 的问题
  • 新层的注册

两个工具：

  • 可视化图: Netron

  • onnx-simplify

onnx 模型快速验证：

  • TensorRT OSS 完成编译之后会生成 trtexec

  • 通过 trtexec 直接加载 onnx， 验证导出模型是否被 TensorRT 支持

  • 根据支持情况，指定算子实现方案

    • 直接基于算子 op 进行开发， 涉及多个 kernel 启动，性能可能相对较差
    • 使用 plugin 进行开发
（3） 自定义组件 TensorRT Plugin

官方文档： https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending

TensorRT 的 Plugin 开发详解
① 观察网络结构， 确认算子的版本、名称、输入、输出、参数及相应权重
② 开发算子的解析模块 parser
③ 继承 PluginV2Ext, 完成 kernel 以及相应 API 的开发
④ 实现 Creator 相关的操作

  • 自定义添加是通过扩展 IPluginV2Ext 和 IPluginCreator 实现的
  • IPluginV2Ext: IPluginV2 的升级版本， 实现自定义插件的基类， 包含版本化和其他格式和单精度的处理
  • IPluginCreator： 自定义层的创建类， 可以通过它获取插件的名称、版本信息、参数等， 也提供网络创建阶段创建插件的方法， 并在推理阶段反序列化它。

4. TensorRT Plugin 开发

自定义组件 TensorRT Plugin 的整体流程如下所示:

  1. 实现 Plugin 类和 PluginCreator 类
  2. 在 InferPlugin.cpp 中注册新 creator
  3. 修改 CMakeLists.txt 并编译 nvinfer_plugin 库
  4. 在 onnx-tensorrt 中注册新的 Plugin：在 onnx 中添加 parser
  5. 编译 nvonnxparse 库
  6. 替换 nvinfer_plugin 和 nvonnxparser 库，使用 trtexec 将 onnx 转换为 tensorrt 模型

下面对每一个步骤进行详细的讲解：

Step1 实现 Plugin 类 和 PluginCreator 类

该步骤可以参考一下 TensorRT 官方 plugin 库： https://github.com/NVIDIA/TensorRT/tree/master/plugin。官方提供了较多的 plugin 插件， 我们可以看到其源码，然后通过模仿源码来学习 plugin 的编写。 这里我们以 nmsPlugin 来看一下自定义的插件如何编写:

.
├── CMakeLists.txt
├── README.md
├── nmsPlugin.cpp  
└── nmsPlugin.h

0 directories, 4 files

这里的 nmsPlugin 主要实现 DetectionOutput 和 NMSPluginCreator 两个类， 前者用于插件的具体实现，后者用于根据需求创建该插件。

class DetectionOutput : public IPluginV2Ext
class NMSPluginCreator : public BaseCreator

DetectionOutput 类继承自 IPluginV2Ext。在阅读其他的插件注册的时候会发现继承自其他类的情况。其实随着 TensorRT 的不断发展，插件接口也在不断地变化，由 v5 版本的IPluginV2Ext，到 v6 版本的 IPluginV2IOExt 和 IPluginV2DynamicExt。 官方的建议是继承自 IPluginV2IOExt 和 IPluginV2DynamicExt。 其实两者在类的编写上类似， 只是其中后者支持动态大小。我们的类需要继承自特定的类，并实现其中的虚函数。

Plugin 类的相关实现函数， 为了方便阅读，我将其分为五类， 分别是准备工作或者结束工作、具体实现和功能函数、序列化问题、插件的基本设置和返回值信息。

（1）准备工作或者结束工作

构造函数: 需要提供三种初始化方式， 分别是通过参数的方式构造函数， 通过 clone 的方式构造函数， 通过序列化的方式构造函数 🌟🌟

// Parameterized constructor
DetectionOutput::DetectionOutput(DetectionOutputParameters params) : param(params){}

// clone Constrcutor
DetectionOutput::DetectionOutput(DetectionOutputParameters params, int C1, int C2, int numPriors)
    : param(params), C1(C1), C2(C2), numPriors(numPriors){}

// constructor
DetectionOutput::DetectionOutput(const void* data, size_t length){
    const char *d = reinterpret_cast<const char*>(data), *a = d;
    param = read<DetectionOutputParameters>(d);
    // Channel size of the locData tensor
    // numPriors * numLocClasses * 4
    C1 = read<int>(d);
    // Channel size of the confData tensor
    // numPriors * param.numClasses
    C2 = read<int>(d);
    // Number of bounding boxes per sample
    numPriors = read<int>(d);
    ASSERT(d == a + length);
}
  • 初始化函数 initialize 和 结束函数 terminate 结束函数。 一般没有什么具体的实现， 返回状态，或者打印 log 而已。
// 在这个插件准备开始 run 之前执行， 做一些初始化工作
// 初始化一些提前开辟空间的参数，一般是 cuda 操作需要的参数
// (例如conv操作需要执行卷积操作，我们就需要提前开辟 weight 和 bias 的显存)，
// 假如我们的算子需要这些参数，则在这里需要提前开辟显存
int DetectionOutput::initialize()
{
    return STATUS_SUCCESS;
}

void DetectionOutput::terminate() {
    gLogVerbose << "NMSPluginDynamic terminate\n";
}
  • clone：将这个 plugin 对象克隆一份 TensorRT 的 builder、network 或者 engine 🌟🌟
// Cloning the plugin
IPluginV2Ext* DetectionOutput::clone() const
{
    // Create a new instance
    IPluginV2Ext* plugin = new DetectionOutput(param, C1, C2, numPriors);

    // Set the namespace
    plugin->setPluginNamespace(mPluginNamespace.c_str());
    return plugin;
}
  • destroy: 一些善后清理工作 
void DetectionOutput::destroy()
{
    delete this;
}

(2) 具体实现和功能函数

  • enqueue： 插件 op 的执行函数 。 其中 enqueue API 实现将具体化的 kernel 调用过程加入执行队列所在流的功能，和具体的 cuda kernel 实现相关联。 🌟🌟
// Plugin layer implementation
int DetectionOutput::enqueue(
    int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)
{
    // Input order {loc, conf, prior}
    const void* const locData = inputs[param.inputOrder[0]];
    const void* const confData = inputs[param.inputOrder[1]];
    const void* const priorData = inputs[param.inputOrder[2]];

    // Output from plugin index 0: topDetections index 1: keepCount
    void* topDetections = outputs[0];
    void* keepCount = outputs[1];

    pluginStatus_t status = detectionInference(stream, batchSize, C1, C2, param.shareLocation,
        param.varianceEncodedInTarget, param.backgroundLabelId, numPriors, param.numClasses, param.topK, param.keepTopK,
        param.confidenceThreshold, param.nmsThreshold, param.codeType, DataType::kFLOAT, locData, priorData,
        DataType::kFLOAT, confData, keepCount, topDetections, workspace, param.isNormalized, param.confSigmoid);
    ASSERT(status == STATUS_SUCCESS);
    return 0;
}
  • attachToContext/detachFromContext：如果这个 op 使用到了一些其他东西， 例如 cublas handle， 可以直接借助 TensorRT 内部提供的 cublas handle。
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
void DetectionOutput::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator){}

// Detach the plugin object from its execution context.
void DetectionOutput::detachFromContext() {}

（3）序列化问题

  • getSerializationSize：序列化时需要写多少字节到 buffer 中 🌟🌟
// Returns the size of serialized parameters
size_t DetectionOutput::getSerializationSize() const
{
    // DetectionOutputParameters, C1, C2, numPriors
    return sizeof(DetectionOutputParameters) + sizeof(int) * 3;
}
  • serialize：把需要用的数据按照顺序序列化到 buffer 里面 🌟🌟
// Serialization of plugin parameters
void DetectionOutput::serialize(void* buffer) const
{
    char *d = reinterpret_cast<char*>(buffer), *a = d;
    write(d, param);
    write(d, C1);
    write(d, C2);
    write(d, numPriors);
    ASSERT(d == a + getSerializationSize());
}

（4）插件的基本设置

getPluginType 和 getPluginVersion 插件的类型或者版本

 // 注意不要重复即可
namespace 
{
    const char* NMS_PLUGIN_VERSION{"1"};
    const char* NMS_PLUGIN_NAME{"NMS_TRT"};
}   // namespace

// Get the plugin type
const char* DetectionOutput::getPluginType() const
{
    return NMS_PLUGIN_NAME;
}

// Get the plugin version
const char* DetectionOutput::getPluginVersion() const
{
    return NMS_PLUGIN_VERSION;
}
  • set/getPluginNamespace：设置/获取插件的命名空间。如果不设置则默认为 “”， 需要注意同一个 namespace 下的 plugin 如果名字相同会冲突。
void DetectionOutput::setPluginNamespace(const char* pluginNamespace)
{
    mPluginNamespace = pluginNamespace;
}

const char* DetectionOutput::getPluginNamespace() const
{
    return mPluginNamespace.c_str();
}
  • supportsFormat：判断格式/数据类型是否合理
// Check if the DataType and Plugin format is supported
bool DetectionOutput::supportsFormat(DataType type, PluginFormat format) const
{
    return (type == DataType::kFLOAT && format == PluginFormat::kNCHW);
}
  • getWorkspaceSize：返回这个插件 op 需要中间显存变量的实际数据大小(bytesize)，这个通过 TensorRT 的接口去获取，是比较规范的方式。 🌟🌟
// Returns the workspace size
size_t DetectionOutput::getWorkspaceSize(int maxBatchSize) const
{
    return detectionInferenceWorkspaceSize(param.shareLocation, maxBatchSize, C1, C2, param.numClasses, numPriors, param.topK, DataType::kFLOAT, DataType::kFLOAT);
}
  • configurePlugin：配置这个插件 op：输入和输出和相关参数的验证和配置。 官方还提到这个配置信息可以告知 TensorRT 去选择合适的算法去调优这个模型 🌟🌟
// Configure the layer with input and output data types.
// inutDims: input Dimensions for the plugin layer
// nInputs : Number of inputs to the plugin layer
// outputDims: output Dimensions from the plugin layer
// nOutputs: number of outputs from the plugin layer
// type: DataType configuration for the plugin layer
// format: format NCHW, NHWC etc
// maxbatchSize: maximum batch size for the plugin layer
void DetectionOutput::configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast, const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize)
{
    ASSERT(nbInputs == 3);
    ASSERT(nbOutputs == 2);

    // Verify all the input dimensions
    for (int i = 0; i < nbInputs; i++)
    {
        ASSERT(inputDims[i].nbDims == 3);
    }

    // Verify all the output dimensions
    for (int i = 0; i < nbOutputs; i++)
    {
        ASSERT(outputDims[i].nbDims == 3);
    }

    // Configure C1, C2 and numPriors
    // Input ordering  C1, C2, numPriors
    C1 = inputDims[param.inputOrder[0]].d[0];
    C2 = inputDims[param.inputOrder[1]].d[0];

    const int nbBoxCoordinates = 4;
    numPriors = inputDims[param.inputOrder[2]].d[1] / nbBoxCoordinates;
    const int numLocClasses = param.shareLocation ? 1 : param.numClasses;

    // Verify C1
    ASSERT(numPriors * numLocClasses * nbBoxCoordinates == inputDims[param.inputOrder[0]].d[0]);

    // Verify C2
    ASSERT(numPriors * param.numClasses == inputDims[param.inputOrder[1]].d[0]);
}

（5） 返回值信息

  • getNbOutputs： 插件 op 返回多少个 Tensor。根据该网络层的实际输出，返回一个整数即可。
//  The nmsPlugin generates an output of shape [batchSize, 1, keepTopK, 7] which contains the same information // as the outputs nmsed box locations, nmsed box scores, and nmsed box class IDs from batchedNMSPlugin, and an // another output of shape [batchSize, 1, 1, 1] which contains the same information as the output nmsed box // count from batchedNMSPlugin.
int DetectionOutput::getNbOutputs() const
{
    return 2;
}
  • getOutputDataType: 返回结果的类型， 一般来说我们插件 op 返回类型与输入类型一致。
// Return the DataType of the plugin output at the requested index.
DataType DetectionOutput::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
{
    // Two outputs
    ASSERT(index == 0 || index == 1); // 这里有两个 output， 需要首先判断 index 是否合理，然后返回类型
    return DataType::kFLOAT;
}
  • getOutputDimensions： TensorRT 支持 Dynamic Shape 的时候， batch 这一维度必须是 explicit 的。 也就是说， TensorRT 处理的维度从以往的三维 [3, -1, -1] 变成了 [1, 3, -1, -1]。最新的 onnx-tensort 也必须设置 explicit 的 bacthsize， 而且这个 bacth 维度在 getOutputDimensions 中是可以获取的。
// Returns output dimensions at given index
Dims DetectionOutput::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
{
    ASSERT(nbInputDims == 3);
    ASSERT(index == 0 || index == 1); // 
    // index 0 : Dimensions 1x param.keepTopK x 7
    // index 1: Dimensions 1x1x1
    if (index == 0)
    {
        return DimsCHW(1, param.keepTopK, 7);
    }
    return DimsCHW(1, 1, 1);
}
  • isOutputBroadcastAcrossBatch/canBroadcastInputAcrossBatch: output 是否进行 boardcast 以及能否进行 boardcast
// Return true if output tensor is broadcast across a batch.
bool DetectionOutput::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const{
    return false;
}

// Return true if plugin can use input that is broadcast across batch without replication.
bool DetectionOutput::canBroadcastInputAcrossBatch(int inputIndex) const{
    return false;
}

PluginCreator 类的相关实现函数，相对来说比较简单。

（1）构造函数

// Plugin creator constructor
NMSPluginCreator::NMSPluginCreator()
{
    // NMS Plugin field meta data {name,  data, type, length}
    mPluginAttributes.emplace_back(PluginField("shareLocation", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("varianceEncodedInTarget", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("backgroundLabelId", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("numClasses", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("topK", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("keepTopK", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("confidenceThreshold", nullptr, PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(PluginField("nmsThreshold", nullptr, PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(PluginField("inputOrder", nullptr, PluginFieldType::kINT32, 3));
    mPluginAttributes.emplace_back(PluginField("confSigmoid", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("isNormalized", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("codeType", nullptr, PluginFieldType::kINT32, 1));

    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

（2）createPlugin：这个成员函数作用是通过 PluginFieldCollection 去创建 plugin， 将 op 需要的权重和参数逐一取出来， 然后调用上文提到的构造函数。

// Creates the NMS plugin
IPluginV2Ext* NMSPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
    const PluginField* fields = fc->fields;
    // Default init values for TF SSD network
    params.codeType = CodeTypeSSD::TF_CENTER;
    params.inputOrder[0] = 0;
    params.inputOrder[1] = 2;
    params.inputOrder[2] = 1;

    // Read configurations from  each fields
    for (int i = 0; i < fc->nbFields; ++i)
    {
        const char* attrName = fields[i].name;
        if (!strcmp(attrName, "shareLocation"))
        {
            ASSERT(fields[i].type == PluginFieldType::kINT32);
            params.shareLocation = static_cast<int>(*(static_cast<const int*>(fields[i].data)));
        }
        // ....
        else if (!strcmp(attrName, "codeType"))
        {
            ASSERT(fields[i].type == PluginFieldType::kINT32);
            params.codeType = static_cast<CodeTypeSSD>(*(static_cast<const int*>(fields[i].data)));
        }
    }

    DetectionOutput* obj = new DetectionOutput(params);
    obj->setPluginNamespace(mNamespace.c_str());
    return obj;
}

（3）deserializePlugin： 这个函数会被 onnx-tensorrt 的一个叫做 TRT_PluginV2 的转换 op 调用， 这个 op 会读取 onnx 模型的 data 数据将其反序列化到 network 中。

IPluginV2Ext* NMSPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
{
    // This object will be deleted when the network is destroyed, which will
    // call NMS::destroy()
    DetectionOutput* obj = new DetectionOutput(serialData, serialLength);
    obj->setPluginNamespace(mNamespace.c_str());
    return obj;
}

（4） 一些插件相关信息配置的函数

// Returns the plugin name
const char* NMSPluginCreator::getPluginName() const
{
    return NMS_PLUGIN_NAME;
}

// Returns the plugin version
const char* NMSPluginCreator::getPluginVersion() const
{
    return NMS_PLUGIN_VERSION;
}

// Returns the plugin field names
const PluginFieldCollection* NMSPluginCreator::getFieldNames()
{
    return &mFC;
}

(5) 获取 PluginFieldCollection。PluginFieldCollection 的主要作用是传递这个插件 op 所需要的权重和参数， 在实际的 engine 推理过程中并不使用， 而在 parse 中会用到（比如 caffe2trt、onnx2trt）。

// Returns the plugin field names
const PluginFieldCollection* NMSPluginCreator::getFieldNames()
{
    return &mFC;
}
Step2： 在 InferPlugin.cpp 中注册新 creator

注册过程， 维护map结构，实现字符串到 creator 的映射。

文件路径：TensorRT/plugin/inferplugin.cpp

将自定义的 plugin 创建器进行初始化。 系统最终通过 creator 映射关系来实现 plugin 的创建和调用

#include "nmsPlugin.h"

extern "C"
{
    bool initLibNvInferPlugins(void* logger, const char* libNamespace)
    {
        initializePlugin<nvinfer1::plugin::NMSPluginCreator>(logger, libNamespace);
        return true;
    }
} // extern "C"
Step3: 修改 CMakeLists.txt 并编译 nvinfer_plugin 库

文件路径：TensorRT/plugin/CMakeLists.txt

// 将 plugin 添加 plugin list
set(PLUGIN_LISTS
    nmsPlugin
    )
setp 4: 在 onnx-tensorrt 中注册新的 Plugin

文件路径： TensorRT/parser/onnx/builtin_op_importers.cpp。 添加方法如下所示：

DEFINE_BUILTIN_OP_IMPORTER(BatchNormalization)  // OP 导入, 参数部分为 OP 名称， 需要与 onnx 一致
{
    // Scale, bias, mean, and variance must be initializers
    // 解析输入信息
    auto scale_weights = inputs.at(1).weights();
    auto bias_weights = inputs.at(2).weights();
    auto mean_weights = inputs.at(3).weights();
    auto variance_weights = inputs.at(4).weights();

    // ...

    // 解析属性信息
    OnnxAttrs attrs(node);
    float eps = attrs.get<float>("epsilon", 1e-5f);

    nvinfer1::Dims dims = tensor_ptr->getDimensions();

    bool need_to_expand_dims = (dims.nbDims == 3);

}
step 5. Plugin 编译支持

确保 TensorRT/Plugin/CMakeLists 保证新增加的 Plugin 能够被正常编译通过。

step 6. 编译 nvonnxparse 库
step 7. 替换 nvinfer_plugin 和 nvonnxparser 库，使用 trtexec 将 onnx 转换为 tensorrt 模型

补充一个学习工具：TensorRT 工具 kaldi-onnx https://github/com/XiaoMi/kaldi-onnx

用于将 kaldi 语音识别工具包神经网络模型移植到onnx模型进行推理的工具。您可以使用 MACE 加快具有高度优化的 neon 内核的 android、ios， linux 或 windows 设备的推断，该工具支持转换 Nnet2 和 Nnet3 模型。

本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议 ，转载请注明出处！