NVIDA-TensorRT部署（一）

TensorRT是一个高性能的深度学习推理(Inference)优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDA的GPU结合起来，能在几乎所有的框架种进行快速和高效的部署推理。

NVIDA开发者指南：开发者指南：： NVIDIA Deep Learning TensorRT 文档

0x01 关于推理（Inference）

0x02 ONNX(Open Neural Network Exchange)

扫描二维码关注公众号，回复： 14585952 查看本文章

0x01 关于推理（Inference）

可从图中看出，训练(training)和推理(inference)的区别：

训练：包含了前向传播和后向传播两个阶段，针对的是训练集。训练时通过误差反向传播来不断修改网络权重(weights)。
推理：只包含了前向传播一个阶段，针对的是除了训练集之外的新数据。可以是测试集，但不完全是，更多的是整个数据集之外的数据，其实就是针对新数据进行预测，预测时，速度是一个很重要的因素。

一般的深度学习项目，训练时为了加快速度，会使用多GPU分布式训练。但在部署推理时，为了降低成本，往往使用单个GPU机器甚至嵌入式平台（NVIDA Jeston）进行部署，部署端也要有与训练时相同的深度学习环境，如caffe、TensorFlow等。

由于训练的网络模型可能会很大，参数很多，而部署的机器性能存在差异，就会导致推理速度慢、延迟高。这对于哪些高实时性的应用场合时很致命的，比如自动驾驶要求实时目标检测、目标追踪等。所以为了提高部署推理的速度，出现了很多轻量级的神经网络，基本做法都是基于现有的经典模型提出一种新的模型结构，然后用这些改造过的模型重新训练，再重新部署。而TensorRT则是对训练好的模型进行优化。TensorRT就只是推理优化器，当你的网络训练完之后，可以将训练模型文件直接丢进TensorRT中，而不再需要依赖深度学习框架。

综上，TensorRT可以看作是一个只有前向传播的深度学习框架，这个框架可以将网络模型进行解析，然后与TensorRT中对应的层进行一一映射，把其他框架的模型统一全部转换到TensorRT中，然后在TensorRT中可以针对NVIDA自家GPU实施优化策略，并进行部署加速。

TensorRT的核心在于对模型算子的优化（合并算子、利用GPU特性选择特定核函数等多种策略），通过TensorRT，能够在Nvidia系列GPU上获得最好的性能。
TensorRT的模型需要在目标GPU上实际运行的方式选择最优算法和配置。
TensorRT生成的模型只能在特定条件下运行（编译trt版本、cuda版本、编译时的GPU型号）

其工作流程可以通过下图得知：

基于TensorRT的发布，也有人在之上做了工作，为每个模型写好应代码，并已写好大量常见的模型代码：GitHub - wang-xinyu/tensorrtx: Implementation of popular deep learning networks with TensorRT network definition API

TensorRT的库文件一览：

0x02 ONNX(Open Neural Network Exchange)

（一）ONNX简介及其可视化

目前TensorRT几乎可以支持所有常用的深度学习框架，对于Caffe和TensorFlow来说，TensorRT可以直接解析它们的网络模型；对于caffe2、pytorch、mxnet、chainer、CNTK等框架则是首先要将模型转为ONNX的通用深度学习模型，然后对ONNX模型做解析。

ONNX是微软和FaceBook携手开发的开放式神经网络交换工具，也就是说不管用什么框架训练，只要转为ONNX模型，就可以放在其他框架上面去inference。这是一种统一的神经网络模型定义和保存方式，上面提到的除了tensorflow之外的其他框架官方应该都对onnx做了支持，而onnx自己开发对tensorflow的支持。

ONNX文件不仅仅存储了神经网络模型的权重，同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其他的辅助信息。可以使用Netron可视化一下：

以yolov5为例

首先生成onnx模型：

pip install onnx

python export.py --weights weights/yolov5s.pt --img 640 --batch 1

生成tflite：

python export.py --weights yolov5s.pt --include tflite --img 320

之后就可以使用netron查看网络了：

pip install netron

在yolov5中执行：

import netron
netron.start("Yolov5s.onnx")

在获得ONNX模型之后，模型部署人员自然就可以将这个模型部署到兼容ONNX的运行环境中，这里一般还会设计到额外的模型转换工作，典型的比如在Android端利用NCNN部署ONNX格式模型，那么就需要ONNX利用NCNN的转换工具转换到NCNN所支持的bin和param格式。

总结：

ONNX的本质，是一种Protobuf格式文件。
Protobuf则通过onnx-ml.proto编译得到onnx-ml.pb.h和onnx-ml.pb.cc或onnx_ml_pb2.py，然后用onnx-ml.pb.cc和代码来操作onnx模型文件，实现增删改。

ModelProto:当加载了一个onnx后，会获得一个ModelProto。它包含一个GraphProto和一些版本，生产者的信息。

GraphProto: 包含了四个repeated数组(可以用来存放N个相同类型的内容，key值为数字序列类型.)。这四个数组分别是node(NodeProto类型)，input(ValueInfoProto类型)，output(ValueInfoProto类型)和initializer(TensorProto类型)；

NodeProto: 存node，放了模型中所有的计算节点,语法结构如下：node表示节点类型，有input属性，是repeated，即重复类型，数组，ouput属性也是同理。name属性是string类型。

ValueInfoProto: 存input，放了模型的输入节点。存output，放了模型中所有的输出节点；

TensorProto: 存initializer，放了模型的所有权重参数。

AttributeProto:每个计算节点中还包含了一个AttributeProto数组，用来描述该节点的属性，比如Conv节点或者说卷积层的属性包含group，pad，strides等等；

可以通过protoc编译onnx-ml.proto，产生onnx-ml.pb.cc文件：

bash make-onnx-pb.sh

onnx-ml.proto则是描述onnx文件如何组成的，具有什么结构，他是操作onnx经常参照的东西。

（二）如何生成ONNX？

首先先组装一个网络，之后将其初始化，最后export为onnx格式的文件，之后就可以在该目录下生成一个demo.onnx：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import os

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(1, 1, 3, padding=1)
        self.relu = nn.ReLU()
        self.conv.weight.data.fill_(1)
        self.conv.bias.data.fill_(0)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

# 这个包对应opset11的导出代码，如果想修改导出的细节，可以在这里修改代码
# import torch.onnx.symbolic_opset11
print("对应opset文件夹代码在这里:", os.path.dirname(torch.onnx.__file__))

model = Model()
dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
    model, 

    # 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
    (dummy,), 

    # 储存的文件路径
    "demo.onnx", 

    # 打印详细信息
    verbose=True, 

    # 为输入和输出节点指定名称，方便后面查看或者操作
    input_names=["image"], 
    output_names=["output"], 

    # 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
    opset_version=11, 

    # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
    # 通常，我们只设置batch为动态，其他的避免动态
    dynamic_axes={
        "image": {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch", 2: "height", 3: "width"},
    }
)

print("Done.!")

（三）如何读取ONNX？

这次我们使用读取上面生成的demo.onnx的文件进行读取，将model中的内容都打印出来：

import onnx
import onnx.helper as helper
import numpy as np

model = onnx.load("demo.change.onnx")

#打印信息
print("==============node信息")
# print(helper.printable_graph(model.graph))
print(model)

conv_weight = model.graph.initializer[0]
conv_bias = model.graph.initializer[1]

# 数据是以protobuf的格式存储的，因此当中的数值会以bytes的类型保存，通过np.frombuffer方法还原成类型为float32的ndarray
print(f"===================={conv_weight.name}==========================")
print(conv_weight.name, np.frombuffer(conv_weight.raw_data, dtype=np.float32))
# 还原出来的ndarray是只读的
print(f"===================={conv_bias.name}==========================")
print(conv_bias.name, np.frombuffer(conv_bias.raw_data, dtype=np.float32))

最后提取出来是这个样子的：

ir_version: 6
producer_name: "pytorch"
producer_version: "1.9"
graph {
  node {
    input: "image"
    input: "conv.weight"
    input: "conv.bias"
    output: "3" // 这里的 3 仅仅是个名字，但是中间过程的node的输入和输出一般仅用数字表示 ouput: "conv_output"但是不简洁
    name: "Conv_0"
    op_type: "Conv"
    attribute {
      name: "dilations"
      ints: 1
      ints: 1
      type: INTS
    }
    attribute {
      name: "group"
      i: 1
      type: INT
    }
    attribute {
      name: "kernel_shape"
      ints: 3
      ints: 3
      type: INTS
    }
    attribute {
      name: "pads"
      ints: 1
      ints: 1
      ints: 1
      ints: 1
      type: INTS
    }
    attribute {
      name: "strides"
      ints: 1
      ints: 1
      type: INTS
    }
  }
  node {
    input: "3"
    output: "output"
    name: "Relu_1"
    op_type: "Relu"
  }
....

model：表示整个onnx的模型，包含图结构和解析器格式、opset版本、导出程序类型。

model.graph：表示图结构，通常是我们netron看到的主要结构。

model.graph.node：表示图中的所有节点，数组，例如conv、bn等节点就是在这里的，通过input、output表示节点之间的连接关系。

model.graph.initializer：权重类的数据大都储存在这里。

model.graph.input：整个模型的输入储存在这里，表明哪个节点是输入节点，shape是多少。

model.graph.output：整个模型的输出储存在这里，表明哪个节点是输出节点，shape是多少。

（四）为什么要编辑ONNX？

假设我们使用PyTorch导出了一个Pt file，之后我们再转换为onnx file，之后再使用Netron进行可视化，onnx是基于protobuf来做数据存储和传输，*.proto后缀文件，其定义是protobuf语法，类似json。如果这个时候我们需要修改一下参数或者其他东西，正常我们会使用PyTorch重新生成一份pt文件，之后再继续上述流程，是比较繁琐的，所以要学会编辑ONNX。ONNX可以看作一个可编辑的记事本。

编辑onnx的demo：

import onnx
import onnx.helper as helper
import numpy as np

model = onnx.load("demo.onnx")

# 可以取出权重
conv_weight = model.graph.initializer[0]
conv_bias = model.graph.initializer[1]
# 修改权
conv_weight.raw_data = np.arange(9, dtype=np.float32).tobytes()

# 修改权重后储存
onnx.save_model(model, "demo.change.onnx")
print("Done.!")

增：一般伴随着增加node和tensor

graph.initializer.append(xxx_tensor)
graph.node.insert(0, xxx_node)

删：

graph.node.remove(xxx_node)

改：

input_node.name = 'data'

（五）生成ONNX

对于一个onnx，他的结构是这样的：

通过import onnx和onnx.helper提供的make_node，make_graph，make_tensor等等接口我们可以轻易的完成一个ONNX模型的构建，需要完成对node，initializer，input，output，graph，model的填充：

根据上图我们可以构建这么一个字典:

graph = helper.make_graph(
    name="mymodel",
    inputs=inputs,
    outputs=outputs,
    nodes=nodes,
    initializer=initializer
)

最后的代码可以如下：

import onnx # pip install onnx>=1.10.2
import onnx.helper as helper
import numpy as np

# https://github.com/onnx/onnx/blob/v1.2.1/onnx/onnx-ml.proto

nodes = [
    helper.make_node(
        name="Conv_0",   # 节点名字，不要和op_type搞混了
        op_type="Conv",  # 节点的算子类型, 比如'Conv'、'Relu'、'Add'这类，详细可以参考onnx给出的算子列表
        inputs=["image", "conv.weight", "conv.bias"],  # 各个输入的名字，结点的输入包含：输入和算子的权重。必有输入X和权重W，偏置B可以作为可选。
        outputs=["3"],  
        pads=[1, 1, 1, 1], # 其他字符串为节点的属性，attributes在官网被明确的给出了，标注了default的属性具备默认值。
        group=1,
        dilations=[1, 1],
        kernel_shape=[3, 3],
        strides=[1, 1]
    ),
    helper.make_node(
        name="ReLU_1",
        op_type="Relu",
        inputs=["3"],
        outputs=["output"]
    )
]

initializer = [
    helper.make_tensor(
        name="conv.weight",
        data_type=helper.TensorProto.DataType.FLOAT,
        dims=[1, 1, 3, 3],
        vals=np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float32).tobytes(),
        raw=True
    ),
    helper.make_tensor(
        name="conv.bias",
        data_type=helper.TensorProto.DataType.FLOAT,
        dims=[1],
        vals=np.array([0.0], dtype=np.float32).tobytes(),
        raw=True
    )
]

inputs = [
    helper.make_value_info(
        name="image",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", 1, 3, 3]
        )
    )
]

outputs = [
    helper.make_value_info(
        name="output",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", 1, 3, 3]
        )
    )
]

graph = helper.make_graph(
    name="mymodel",
    inputs=inputs,
    outputs=outputs,
    nodes=nodes,
    initializer=initializer
)

# 如果名字不是ai.onnx，netron解析就不是太一样了
opset = [
    helper.make_operatorsetid("ai.onnx", 11)
]

# producer主要是保持和pytorch一致
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.9")
onnx.save_model(model, "my.onnx")

print(model)
print("Done.!")

（六）正确导出ONNX

对于任何用到shape、size返回值的参数时，例如：tensor.view(tensor.size(0), -1)这类操作，避免直接使用tensor.size的返回值，而是加上int转换，tensor.view(int(tensor.size(0)), -1)，断开跟踪。

对于nn.Upsample或nn.functional.interpolate函数，使用scale_factor指定倍率，而不是使用size参数指定大小。
对于reshape、view操作时，-1的指定请放到batch维度。其他维度可以计算出来即可。batch维度禁止指定为大于-1的明确数字。
torch.onnx.export指定dynamic_axes参数，并且只指定batch维度，禁止其他动态。
使用opset_version=11，不要低于11。

（七）ONNX解析器

onnx解析器有两个选项，libnvonnxparser.so或者GitHub - onnx/onnx-tensorrt: ONNX-TensorRT: TensorRT backend for ONNX（源代码）。

如果是去源代码下下载解析器的，需要匹配自己的tensorRT版本。

源代码中主要关注的是 builtin_op_importers.cpp 文件，这个文件中定义了各个算子是怎样从 ONNX 算子解析成 TensorRT 算子的，如果我们后续有新的算子，也可以自己修改该源文件来支持新算子。

0x03 Hello TensorRT

首先使用TensorRT构建一个全连接网络。

TensorRT的工作流程如下：

首先定义网络。
优化builder参数。
通过builder生成engine，用于模型保存、推理等。
engine可以通过序列化和逆序列化转化模型数据类型（转化为二进制byte文件，加快传输速率），再进一步推动模型由输入张量到输出张量的推理。

构建网络：

代码：

// tensorRT include
#include <NvInfer.h>
#include <NvInferRuntime.h>

// cuda include
#include <cuda_runtime.h>

// system include
#include <stdio.h>

// 创建日志，用于打印信息
class TRTLogger : public nvinfer1::ILogger{
public:
    virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{
        if(severity <= Severity::kVERBOSE){
            printf("%d: %s\n", severity, msg);
        }
    }
};

nvinfer1::Weights make_weights(float* ptr, int n){
    nvinfer1::Weights w;
    w.count = n;
    w.type = nvinfer1::DataType::kFLOAT;
    w.values = ptr;
    return w;
}

int main(){
     
    TRTLogger logger; // logger是必要的，用来捕捉warning和info等

    // ----------------------------- 1. 定义 builder, config 和network -----------------------------
    // 这是基本需要的组件
    //形象的理解是你需要一个builder去build这个网络，网络自身有结构，这个结构可以有不同的配置
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
    // 创建一个构建配置，指定TensorRT应该如何优化模型，tensorRT生成的模型只能在特定配置下运行
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    // 创建网络定义，其中createNetworkV2(1)表示采用显性batch size，新版tensorRT(>=7.0)时，不建议采用0非显性batch size
    // 因此贯穿以后，请都采用createNetworkV2(1)而非createNetworkV2(0)或者createNetwork
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);

    // 构建一个模型
    /*
        Network definition:

        image
          |
        linear (fully connected)  input = 3, output = 2, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]
          |
        sigmoid
          |
        prob
    */

    // ----------------------------- 2. 输入，模型结构和输出的基本信息 -----------------------------
    const int num_input = 3;   // in_channel
    const int num_output = 2;  // out_channel
    float layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5}; // 前3个给w1的rgb，后3个给w2的rgb 
    float layer1_bias_values[]   = {0.3, 0.8};

    //输入指定数据的名称、数据类型和完整维度，将输入层添加到网络
    nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(1, num_input, 1, 1));
    nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 6);
    nvinfer1::Weights layer1_bias   = make_weights(layer1_bias_values, 2);
    //添加全连接层
    auto layer1 = network->addFullyConnected(*input, num_output, layer1_weight, layer1_bias);      // 注意对input进行了解引用
    //添加激活层 
    auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); // 注意更严谨的写法是*(layer1->getOutput(0)) 即对getOutput返回的指针进行解引用
    
    // 将我们需要的prob标记为输出
    network->markOutput(*prob->getOutput(0));

    printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f); // 256Mib
    config->setMaxWorkspaceSize(1 << 28);
    builder->setMaxBatchSize(1); // 推理时 batchSize = 1 

    // ----------------------------- 3. 生成engine模型文件 -----------------------------
    //TensorRT 7.1.0版本已弃用buildCudaEngine方法，统一使用buildEngineWithConfig方法
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if(engine == nullptr){
        printf("Build engine failed.\n");
        return -1;
    }

    // ----------------------------- 4. 序列化模型文件并存储 -----------------------------
    // 将模型序列化，并储存为文件
    nvinfer1::IHostMemory* model_data = engine->serialize();
    FILE* f = fopen("engine.trtmodel", "wb");
    fwrite(model_data->data(), 1, model_data->size(), f);
    fclose(f);

    // 卸载顺序按照构建顺序倒序
    model_data->destroy();
    engine->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
    printf("Done.\n");
    return 0;
}

注意：

必须使用createNetworkV2，并指定为1（表示显性batch）。createNetwork已经废弃，非显性batch官方不推荐。这个方式直接影响推理时enqueue还是enqueueV2。
builder、config等指针，记得释放，否则会有内存泄漏，使用ptr->destroy()释放。
markOutput表示是该模型的输出节点，mark几次，就有几个输出，addInput几次就有几个输入。这与推理时相呼应。
workspaceSize是工作空间大小，某些layer需要使用额外存储时，不会自己分配空间，而是为了内存复用，直接找tensorRT要workspace空间。
一定要记住，保存的模型只能适配编译时的trt版本、编译时指定的设备。也只能保证在这种配置下是最优的。如果用trt跨不同设备执行，有时候可以运行，但不是最优的，也不推荐。

0x04 Hello Inference

使用上面构建的全连接层网络来进行推理：

void inference(){

    // ------------------------------ 1. 准备模型并加载   ----------------------------
    TRTLogger logger;
    auto engine_data = load_file("engine.trtmodel");
    // 执行推理前，需要创建一个推理的runtime接口实例。与builer一样，runtime需要logger：
    nvinfer1::IRuntime* runtime   = nvinfer1::createInferRuntime(logger);
    // 将模型从读取到engine_data中，则可以对其进行反序列化以获得engine
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
    if(engine == nullptr){
        printf("Deserialize cuda engine failed.\n");
        runtime->destroy();
        return;
    }
	// 创建执行上下文
    nvinfer1::IExecutionContext* execution_context = engine->createExecutionContext();
	// 异步
    cudaStream_t stream = nullptr;
    // 创建CUDA流，以确定这个batch的推理是独立的
    cudaStreamCreate(&stream);

    /*
        Network definition:

        image
          |
        linear (fully connected)  input = 3, output = 2, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]
          |
        sigmoid
          |
        prob
    */

    // ------------------------------ 2. 准备好要推理的数据并搬运到GPU   ----------------------------
    float input_data_host[] = {1, 2, 3};
    float* input_data_device = nullptr;

    float output_data_host[2];
    float* output_data_device = nullptr;
    cudaMalloc(&input_data_device, sizeof(input_data_host));
    cudaMalloc(&output_data_device, sizeof(output_data_host));
    cudaMemcpyAsync(input_data_device, input_data_host, sizeof(input_data_host), cudaMemcpyHostToDevice, stream);
    // 用一个指针数组指定input和output在gpu中的指针。
    float* bindings[] = {input_data_device, output_data_device};

    // ------------------------------ 3. 推理并将结果搬运回CPU   ----------------------------
	// 送去推理
    bool success      = execution_context->enqueueV2((void**)bindings, stream, nullptr);
    cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);

    printf("output_data_host = %f, %f\n", output_data_host[0], output_data_host[1]);

    // ------------------------------ 4. 释放内存 ----------------------------
    printf("Clean memory\n");
    cudaStreamDestroy(stream);
    execution_context->destroy();
    engine->destroy();
    runtime->destroy();

    // ------------------------------ 5. 手动推理进行验证 ----------------------------
    const int num_input = 3;
    const int num_output = 2;
    float layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5};
    float layer1_bias_values[]   = {0.3, 0.8};

    printf("手动验证计算结果：\n");
    for(int io = 0; io < num_output; ++io){
        float output_host = layer1_bias_values[io];
        for(int ii = 0; ii < num_input; ++ii){
            output_host += layer1_weight_values[io * num_input + ii] * input_data_host[ii];
        }

        // sigmoid
        float prob = 1 / (1 + exp(-output_host));
        printf("output_prob[%d] = %f\n", io, prob);
    }
}

需要注意的是：

bindings是tensorRT对输入输出张量的描述，bindings = input-tensor + output-tensor。比如input有a，output有b, c, d，那么bindings = [a, b, c, d]，bindings[0] = a，bindings[2] = c。此时看到engine->getBindingDimensions(0)你得知道获取的是什么。
enqueueV2是异步推理，加入到stream队列等待执行。输入的bindings则是tensors的指针（注意是device pointer）。其shape对应于编译时指定的输入输出的shape（这里只演示全部shape静态）。
createExecutionContext可以执行多次，允许一个引擎具有多个执行上下文，不过看看就好，别当真。

总结以上，推理的步骤如下：

准备模型并加载。
创建runtime：createInferRuntime(logger)
使用运行时，步骤如下：
- 反序列化创建engine, 得为engine提供数据：runtime->deserializeCudaEngine(modelData, modelSize),其中modelData包含的是input和output的名字，形状，大小和数据类型。
- 从engine创建执行上下文:engine->createExecutionContext()
创建CUDA流cudaStreamCreate(&stream)：
- CUDA编程流是组织异步工作的一种方式，创建流来确定batch推理的独立。
- 为每个独立的batch使用IExecutionContext，并为每个独立批次使用cudaStreamCreate创建CUDA流。
数据准备：
- 在host上声明input数据和output数组大小，搬运到GPU上。
- 要执行inference，必须用一个指针数组指定input和output在gpu中的指针。
- 推理并将output搬运回CPU
启动所有工作后，与所有流同步以等待结果:cudaStreamSynchronize
按照与创建相反顺序释放内存。

0x05 Dynamic shapes

Dynamic shapes指的是我们可以在runtime推理阶段来指定some或者all输入数据的维度。一般需要指定为dynamic的是batch_size这一个维度，使得我们可以根据自己实际情况动态设置batch，而不需要每次都重新生成engine文件。

动态shape，即编译时指定可动态的范围[L-H]，推理时可以允许L<=shaple<=H

如何生成及使用Dynamic shapes的engine的大致步骤如下：

使用最新的接口创建NetworkDefinition对象：也就是上面提及到的createNetworkV2()。
对于Input Tensor中的Dynamic的维度，通过-1来占位这个维度。
在build阶段设置一个或多个optimization profiles，用来指定在runtime阶段inputs允许的维度范围，一般设置三个profiles：最小、最大和最优。之后通过上述的设置，就可以在build阶段生成一个带Dynamic shapes的engine文件，然后就是在推理阶段如何使用这个engine。
如何使用带Dynamic shapes的engine：
- 创建一个execution context，此时的context也是Dynamic shapes的。
- 将想要设置的input dimension绑定context，那么此时context处理的维度确定了。
- 基于这个context进行推理。

那么首先看看代码：

我们要构建一个如下的模型：

/*
    Network definition:

    image
      |
    conv(3x3, pad=1)  input = 1, output = 1, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5], [0.2, 0.2, 0.1]], b=0.0
      |
    relu
      |
    prob
*/

之后我们进行输入，模型结构和输出的基本信息，需要注意的是，我们必须让必须让NetworkDefinition的维度定义为-1

// 如果要使用动态shape，必须让NetworkDefinition的维度定义为-1，in_channel是固定的
nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(-1, num_input, -1, -1));

最关键的是去设置这个profile，这个是动态shape的重点：

	// 如果模型有多个输入，则必须多个profile
    auto profile = builder->createOptimizationProfile();

    // 配置最小允许1 x 1 x 3 x 3
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, num_input, 3, 3));
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, num_input, 3, 3));

    // 配置最大允许10 x 1 x 5 x 5
    // if networkDims.d[i] != -1, then minDims.d[i] == optDims.d[i] == maxDims.d[i] == networkDims.d[i]
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(maxBatchSize, num_input, 5, 5));
    config->addOptimizationProfile(profile);

那么整个程序如下：

// tensorRT include
#include <NvInfer.h>
#include <NvInferRuntime.h>

// cuda include
#include <cuda_runtime.h>

// system include
#include <stdio.h>
#include <math.h>

#include <iostream> 
#include <fstream> // 后面要用到ios这个库
#include <vector>

using namespace std;

class TRTLogger : public nvinfer1::ILogger{
public:
    virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{
        if(severity <= Severity::kINFO){
            printf("%d: %s\n", severity, msg);
        }
    }
} logger;

nvinfer1::Weights make_weights(float* ptr, int n){
    nvinfer1::Weights w;
    w.count = n;
    w.type = nvinfer1::DataType::kFLOAT;
    w.values = ptr;
    return w;
}

bool build_model(){
    TRTLogger logger;

    // ----------------------------- 1. 定义 builder, config 和network -----------------------------
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);

    // 构建一个模型
    /*
        Network definition:

        image
          |
        conv(3x3, pad=1)  input = 1, output = 1, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5], [0.2, 0.2, 0.1]], b=0.0
          |
        relu
          |
        prob
    */


    // ----------------------------- 2. 输入，模型结构和输出的基本信息 -----------------------------
    const int num_input = 1;
    const int num_output = 1;
    float layer1_weight_values[] = {
        1.0, 2.0, 3.1, 
        0.1, 0.1, 0.1, 
        0.2, 0.2, 0.2
    }; // 行优先
    float layer1_bias_values[]   = {0.0};

    // 如果要使用动态shape，必须让NetworkDefinition的维度定义为-1，in_channel是固定的
    nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(-1, num_input, -1, -1));
    nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 9);
    nvinfer1::Weights layer1_bias   = make_weights(layer1_bias_values, 1);
    auto layer1 = network->addConvolution(*input, num_output, nvinfer1::DimsHW(3, 3), layer1_weight, layer1_bias);
    layer1->setPadding(nvinfer1::DimsHW(1, 1));

    auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kRELU); // *(layer1->getOutput(0))
     
    // 将我们需要的prob标记为输出
    network->markOutput(*prob->getOutput(0));

    int maxBatchSize = 10;
    printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
    // 配置暂存存储器，用于layer实现的临时存储，也用于保存中间激活值
    config->setMaxWorkspaceSize(1 << 28);

    // --------------------------------- 2.1 关于profile ----------------------------------
    // 如果模型有多个输入，则必须多个profile
    auto profile = builder->createOptimizationProfile();

    // 配置最小允许1 x 1 x 3 x 3
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, num_input, 3, 3));
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, num_input, 3, 3));

    // 配置最大允许10 x 1 x 5 x 5
    // if networkDims.d[i] != -1, then minDims.d[i] == optDims.d[i] == maxDims.d[i] == networkDims.d[i]
    profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(maxBatchSize, num_input, 5, 5));
    config->addOptimizationProfile(profile);

    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if(engine == nullptr){
        printf("Build engine failed.\n");
        return false;
    }

    // -------------------------- 3. 序列化 ----------------------------------
    // 将模型序列化，并储存为文件
    nvinfer1::IHostMemory* model_data = engine->serialize();
    FILE* f = fopen("engine.trtmodel", "wb");
    fwrite(model_data->data(), 1, model_data->size(), f);
    fclose(f);

    // 卸载顺序按照构建顺序倒序
    model_data->destroy();
    engine->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
    printf("Done.\n");
    return true;
}

vector<unsigned char> load_file(const string& file){
    ifstream in(file, ios::in | ios::binary);
    if (!in.is_open())
        return {};

    in.seekg(0, ios::end);
    size_t length = in.tellg();

    std::vector<uint8_t> data;
    if (length > 0){
        in.seekg(0, ios::beg);
        data.resize(length);

        in.read((char*)&data[0], length);
    }
    in.close();
    return data;
}

void inference(){
    // ------------------------------- 1. 加载model并反序列化 -------------------------------
    TRTLogger logger;
    auto engine_data = load_file("engine.trtmodel");
    nvinfer1::IRuntime* runtime   = nvinfer1::createInferRuntime(logger);
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
    if(engine == nullptr){
        printf("Deserialize cuda engine failed.\n");
        runtime->destroy();
        return;
    }

    nvinfer1::IExecutionContext* execution_context = engine->createExecutionContext();
    cudaStream_t stream = nullptr;
    cudaStreamCreate(&stream);

    /*
        Network definition:

        image
          |
        conv(3x3, pad=1)  input = 1, output = 1, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5], [0.2, 0.2, 0.1]], b=0.0
          |
        relu
          |
        prob
    */

    // ------------------------------- 2. 输入与输出 -------------------------------
    float input_data_host[] = {
        // batch 0
        1,   1,   1,
        1,   1,   1,
        1,   1,   1,

        // batch 1
        -1,   1,   1,
        1,   0,   1,
        1,   1,   -1
    };
    float* input_data_device = nullptr;

    // 3x3输入，对应3x3输出
    int ib = 2;
    int iw = 3;
    int ih = 3;
    float output_data_host[ib * iw * ih];
    float* output_data_device = nullptr;
    cudaMalloc(&input_data_device, sizeof(input_data_host));
    cudaMalloc(&output_data_device, sizeof(output_data_host));
    cudaMemcpyAsync(input_data_device, input_data_host, sizeof(input_data_host), cudaMemcpyHostToDevice, stream);


    // ------------------------------- 3. 推理 -------------------------------
    // 明确当前推理时，使用的数据输入大小
    execution_context->setBindingDimensions(0, nvinfer1::Dims4(ib, 1, ih, iw));
    float* bindings[] = {input_data_device, output_data_device};
    bool success      = execution_context->enqueueV2((void**)bindings, stream, nullptr);
    cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);


    // ------------------------------- 4. 输出结果 -------------------------------
    for(int b = 0; b < ib; ++b){
        printf("batch %d. output_data_host = \n", b);
        for(int i = 0; i < iw * ih; ++i){
            printf("%f, ", output_data_host[b * iw * ih + i]);
            if((i + 1) % iw == 0)
                printf("\n");
        }
    }

    printf("Clean memory\n");
    cudaStreamDestroy(stream);
    cudaFree(input_data_device);
    cudaFree(output_data_device);
    execution_context->destroy();
    engine->destroy();
    runtime->destroy();
}

int main(){

    if(!build_model()){
        return -1;
    }
    inference();
    return 0;
}

那么总结来说：

（一）构建网络时

必须在模型定义时，输入维度给定为-1，否则该维度不会动态。如果为onnx文件，则onnx文件打开后应该看到动态或者-1；如果模型中存在reshape类操作，那么reshape的参数必须随动态进行计算，而大部分时候都是问题。除非为全卷积模型，否则大部门时候只需要batch_size维度设置为动态，其他维度尽量避免设置动态。
配置profile：

create: builder->createOptimizationProfile()。

set: setDimensions()设置kMIN, kOPT, kMAX的一系列输入尺寸范围。

add:config->addOptimizationProfile(profile);添加profile到网络配置中。

（二）推理阶段

需要在选择profile的索引后设置input维度：execution_context->setBindingDimensions(0, nvinfer1::Dims4(1, 1, 3, 3));
在运行时，向engine请求绑定维度会返回用于构建网络的相同维度。这意味着，得到的还是动态的维度[-1, in_channel, -1, -1]：

engine.getBindingDimensions(0) // return [-1, 1, -1, -1]

获取当前的实际维度，需要查询执行上下文：

context.getBindingDimensions(0) // return [3, 1, 3, 3]

0x06 插件实现

TensorRT插件的存在目的，主要是为了让我们实现TensorRT目前还不支持的算子，所以我们需要TensorRT的plugin去实现我们自己的算子。

TensorRT的官方plugin库长啥样嘞：TensorRT/plugin at master · NVIDIA/TensorRT · GitHub

如果要添加自己的算子，可以在官方的plugin库里头进行修改添加，然后编译官方的plugin库。将新生成的libnvinfer_plugin.so替换原本的.so文件即可。或者自己写一个类似于官方plugin的组件，将名称替换一下，同样生成.so，在TensorRT的推理项目中引用这个动态链接库即可。

首先，我们先生成带着我们自己算子的ONNX文件，我们将这个算子称为MYSELU：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import torch.autograd
import os

class MYSELUImpl(torch.autograd.Function):

    # reference: https://pytorch.org/docs/1.10/onnx.html#torch-autograd-functions
    @staticmethod
    def symbolic(g, x, p):
        print("==================================call symbolic")
        return g.op("MYSELU", x, p, 
            g.op("Constant", value_t=torch.tensor([3, 2, 1], dtype=torch.float32)),
            attr1_s="这是字符串属性", 
            attr2_i=[1, 2, 3], 
            attr3_f=222
        )

    @staticmethod
    def forward(ctx, x, p):
        return x * 1 / (1 + torch.exp(-x))


class MYSELU(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.param = nn.parameter.Parameter(torch.arange(n).float())

    def forward(self, x):
        return MYSELUImpl.apply(x, self.param)


class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(1, 1, 3, padding=1)
        self.myselu = MYSELU(3)
        self.conv.weight.data.fill_(1)
        self.conv.bias.data.fill_(0)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.myselu(x)
        return x


# 这个包对应opset11的导出代码，如果想修改导出的细节，可以在这里修改代码
# import torch.onnx.symbolic_opset11
print("对应opset文件夹代码在这里：", os.path.dirname(torch.onnx.__file__))

model = Model().eval()
input = torch.tensor([
    # batch 0
    [
        [1,   1,   1],
        [1,   1,   1],
        [1,   1,   1],
    ],
        # batch 1
    [
        [-1,   1,   1],
        [1,   0,   1],
        [1,   1,   -1]
    ]
], dtype=torch.float32).view(2, 1, 3, 3)

output = model(input)
print(f"inference output = \n{output}")

dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
    model, 

    # 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
    (dummy,), 

    # 储存的文件路径
    "workspace/demo.onnx", 

    # 打印详细信息
    verbose=True, 

    # 为输入和输出节点指定名称，方便后面查看或者操作
    input_names=["image"], 
    output_names=["output"], 

    # 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
    opset_version=11, 

    # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
    # 通常，我们只设置batch为动态，其他的避免动态
    dynamic_axes={
        "image": {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch", 2: "height", 3: "width"},
    },

    # 对于插件，需要禁用onnx检查
    enable_onnx_checker=False
)

print("Done.!")

（一）按照官方的排布方式编写插件

准备一个自己的插件：my-plugin.cpp和my-plugin.hpp，然后ctrl c v官方代码，名字换成自己的，以最新的IPluginV2DynamicExt类为接口。

我们需要写两个类：

class MySELUPlugin : public IPluginV2DynamicExt，继承了IPluginV2DynamicExt，是插件类，用于写插件具体的实现。
class MySELUPluginCreator : public IPluginCreator，是插件工厂类，用于根据需求创建该插件。

MySELUPlugin插件类

可以先看看hpp中的声明：

class MySELUPlugin : public IPluginV2DynamicExt
{
public:
    // 两种参数构造方式
    MySELUPlugin(const std::string name, const std::string attr1, float attr3);

    MySELUPlugin(const std::string name, const void* data, size_t length);

    // It doesn't make sense to make MySELUPlugin without arguments, so we delete default constructor.
    MySELUPlugin() = delete;

    // 插件输出数量
    int getNbOutputs() const noexcept override;

    // 插件输出类型是什么
    virtual nvinfer1::DataType getOutputDataType(
        int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override{
        return inputTypes[0];
    }
    // 输入为动态shape的情况 获取当前shape
    virtual nvinfer1::DimsExprs getOutputDimensions(
        	int32_t outputIndex, const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override;

    int initialize() noexcept override;

    void terminate() noexcept override;
    // 使用内存/显存，实现内存复用
    virtual size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int32_t nbInputs, const nvinfer1::PluginTensorDesc* outputs,
        	int32_t nbOutputs) const noexcept override{
        return 0;
    }
    
    int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
            const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
    // 获取序列化的大小
    size_t getSerializationSize() const noexcept override;

    void serialize(void* buffer) const noexcept override;
    // 选择配置
    virtual void configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs,
        const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept;

    virtual bool supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;

    const char* getPluginType() const noexcept override;

    const char* getPluginVersion() const noexcept override;

    void destroy() noexcept override;

    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;

    void setPluginNamespace(const char* pluginNamespace) noexcept override;

    const char* getPluginNamespace() const noexcept override;

private:
    const std::string mLayerName;
    std::string mattr1;
    float mattr3;
    size_t mInputVolume;
    std::string mNamespace;
};

之后构造函数有两个，一个用于推理阶段、一个用于编译阶段：

// 定义插件类MYSELUPlugin
MySELUPlugin::MySELUPlugin(const std::string name, const std::string attr1, float attr3)
    : mLayerName(name),
    mattr1(attr1),
    mattr3(attr3)
{
    printf("==================== 编译阶段，attr1 = %s, attr3 = %f\n", attr1.c_str(), attr3);
}

MySELUPlugin::MySELUPlugin(const std::string name, const void* data, size_t length)
    : mLayerName(name)
{
    // Deserialize in the same order as serialization
    const char* d = static_cast<const char*>(data);
    const char* a = d;

    int nstr = readFromBuffer<int>(d);
    mattr1 = std::string(d, d + nstr);

    d += nstr;
    mattr3 = readFromBuffer<float>(d);
    assert(d == (a + length));

    printf("==================== 推理阶段，attr1 = %s, attr3 = %f\n", mattr1.c_str(), mattr3);
}

注意我们需要把默认的构造函数删掉：

MySELUPlugin() = delete;

getNbOutputs：插件OP返回多少个Tensor，这个MySELUPlugin的操作只输出一个Tensor，也就是一个Output，所以直接return 1：

int MySELUPlugin::getNbOutputs() const noexcept
{
    return 1;
}

initialize：主要初始化一些提前开辟空间的参数，一般是一些cuda操作需要的参数(例如conv操作需要执行卷积操作，我们就需要提前开辟weight和bias的显存)，假如我们的算子需要这些参数，则在这里需要提前开辟显存。

需要注意的是，如果插件算子需要开辟比较大的显存空间，不建议自己去申请显存空间，可以使用Tensorrt官方接口传过来的workspace指针来获取显存空间。因为如果这个插件被一个网络调用了很多次，而这个插件op需要开辟很多显存空间，那么TensorRT在构建network的时候会根据这个插件被调用的次数开辟很多显存，很容易导致显存溢出。

getOutputDataType：返回结果的类型，一般来说我们插件返回类型与输入类型是一致的：

virtual nvinfer1::DataType getOutputDataType(
    int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override{
    return inputTypes[0];
}

getWorkspaceSize：这个函数需要返回这个插件OP需要中间显存变量的实际数据大小（bytesize），这个是通过TensorRT的接口去获取，是比较规范的方式。我们需要在这里确定这个OP需要多大的显存空间去运行，在实际运行的时候就可以直接使用TensorRT开辟好的空间而不是自己去申请显存空间。

virtual size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int32_t nbInputs, const nvinfer1::PluginTensorDesc* outputs,
        	int32_t nbOutputs) const noexcept override{
        return 0;
    }

enqueue：实际插件OP的执行函数，CUDA的操作就是在这里：

int MySELUPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
    const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
    void* output = outputs[0];
    size_t volume = 1;
    for (int i = 0; i < inputDesc->dims.nbDims; i++){
        volume *= inputDesc->dims.d[i];
    }
    mInputVolume = volume;

    myselu_inference(
        static_cast<const float*>(inputs[0]), 
        static_cast<float*>(output), 
        mInputVolume,
        stream
    );
    return 0;
}

在其中我们使用的是myselu_inference函数，这个函数是在.cu中声明的，也就是cuda函数：

void myselu_inference(const float* x, float* output, int n, cudaStream_t stream){

    const int nthreads = 512;
    int block_size = n < nthreads ? n : nthreads;
    int grid_size = (n + block_size - 1) / block_size;
    myselu_kernel<<<grid_size, block_size, 0, stream>>>(x, output, n);
}

getOutputDimensions：TensorRT支持动态shape，在这里获取该层的输出维度是多少，这里不改变输入尺寸，故与输入尺寸相同。

nvinfer1::DimsExprs MySELUPlugin::getOutputDimensions(int32_t outputIndex, const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
{
    // MySELUping不改变输入尺寸，所以输出尺寸将与输入尺寸相同
    return *inputs;
}

configurePlugin：配置插件格式：告诉你目前这个层所采用的数据格式和类型。

void MySELUPlugin::configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs,
    const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept{
    // Validate input arguments

    auto type = in->desc.type;
    auto format = in->desc.format;
    assert(nbOutputs == 1);
    assert(type == DataType::kFLOAT);
    assert(format == PluginFormat::kLINEAR);
}

clone:将这个plugin对象克隆一份给TensorRT的builder、network或者engine,clone成员函数主要用于传递不变的权重和参数，将plugin复制n多份，从而可以被不同engine或者builder或者network使用。

IPluginV2DynamicExt* MySELUPlugin::clone() const noexcept
{
    printf("===================克隆插件=================\n");
    auto plugin = new MySELUPlugin(mLayerName, mattr1, mattr3);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

这个函数调用了上述的构造函数。

MySELUPluginCreator插件工厂类：

class MySELUPluginCreator : public IPluginCreator
{
public:
    MySELUPluginCreator();

    const char* getPluginName() const noexcept override;

    const char* getPluginVersion() const noexcept override;

    const PluginFieldCollection* getFieldNames() noexcept override;

    IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override;

    IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;

    void setPluginNamespace(const char* pluginNamespace) noexcept override;

    const char* getPluginNamespace() const noexcept override;

private:
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
    std::string mNamespace;
};

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

// 创建plugin
IPluginV2* MySELUPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
{
    std::string attr1;
    float attr3;
    const PluginField* fields = fc->fields;

    // Parse fields from PluginFieldCollection
    for (int i = 0; i < fc->nbFields; i++)
    {
        if (strcmp(fields[i].name, "attr1") == 0)
        {
            assert(fields[i].type == PluginFieldType::kCHAR);
            auto cp = static_cast<const char*>(fields[i].data);
            attr1 = std::string(cp, cp + fields[i].length);
        }
        else if (strcmp(fields[i].name, "attr3") == 0)
        {
            assert(fields[i].type == PluginFieldType::kFLOAT32);
            attr3 = *(static_cast<const float*>(fields[i].data));
        }
    }
    return new MySELUPlugin(name, attr1, attr3);
}

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

// 反序列化插件参数进行创建
IPluginV2* MySELUPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
{
    // This object will be deleted when the network is destroyed, which will
    // call MySELUPlugin::destroy()
    return new MySELUPlugin(name, serialData, serialLength);
}

（二）总结——实现插件的要点

导出onnx时，需要为module增加symbolic函数。

参照：torch.onnx — PyTorch 1.10 文档

g.op对应的名称需要与下面解析器的名称对应。
需要在解析器下添加对插件op的解析：在文件builtin_op_importers.cpp下添加。这个文件是我们添加插件时需要编写的文件。

DEFINE_BUILTIN_OP_IMPORTER(MYSELU)
{
    printf("\033[31m=======================call MYSELU==============\033[0m\n");
    OnnxAttrs attrs(node, ctx);
    const std::string pluginName{node.op_type()};
    const std::string pluginVersion{attrs.get<std::string>("plugin_version", "1")};
    const std::string pluginNamespace{attrs.get<std::string>("plugin_namespace", "")};

    LOG_INFO("Searching for plugin: " << pluginName << ", plugin_version: " << pluginVersion << ", plugin_namespace: " << pluginNamespace);
    nvinfer1::IPluginCreator* creator = importPluginCreator(pluginName, pluginVersion, pluginNamespace);
    ASSERT(creator && "Plugin not found, are the plugin name, version, and namespace correct?", ErrorCode::kUNSUPPORTED_NODE);

    const nvinfer1::PluginFieldCollection* fieldNames = creator->getFieldNames();
    // Field data needs to be type erased, we use fieldData for temporary allocations.
    string_map<std::vector<uint8_t>> fieldData{};
    std::vector<nvinfer1::PluginField> fields = loadFields(fieldData, attrs, fieldNames, ctx);

    const auto plugin = createPlugin(getNodeName(node), creator, fields);
    ASSERT(plugin && "Could not create plugin", ErrorCode::kUNSUPPORTED_NODE);

    std::vector<nvinfer1::ITensor*> pluginInputs{};
    for (auto& input : inputs)
    {
        pluginInputs.emplace_back(&convertToTensor(input, ctx));
    }
    LOG_INFO("Successfully created plugin: " << pluginName);
    auto* layer = ctx->network()->addPluginV2(pluginInputs.data(), pluginInputs.size(), *plugin);
    ctx->registerLayer(layer, getNodeName(node));
    RETURN_ALL_OUTPUTS(layer);
}

这个家伙是个宏定义，原型在这里：

#define DEFINE_BUILTIN_OP_IMPORTER(op)                                                                                 \
    NodeImportResult import##op(                                                                                       \
        IImporterContext* ctx, ::onnx::NodeProto const& node, std::vector<TensorOrWeights>& inputs);         \
    static const bool op##_registered_builtin_op = registerBuiltinOpImporter(#op, import##op);                         \
    IGNORE_UNUSED_GLOBAL(op##_registered_builtin_op);                                                                  \
    NodeImportResult import##op(                                                                                       \
        IImporterContext* ctx, ::onnx::NodeProto const& node, std::vector<TensorOrWeights>& inputs)

创建MySELUPluginCreator，插件创建器

// 静态类字段的初始化
PluginFieldCollection MySELUPluginCreator::mFC{}; // FieldCollection 字段收集
std::vector<PluginField> MySELUPluginCreator::mPluginAttributes;
// 实际注册时，注册的是创建器，交给tensorRT管理
REGISTER_TENSORRT_PLUGIN(MySELUPluginCreator);

定义插件类MySELUPlugin，Creator创建器来实例化MySELUPlugin类。

总结使用流程：

编译阶段
- 通过MySELUPluginCreator::createPlugin创建plugin。
- 期间会调用MySELUPlugin::clone克隆插件
- 调用MySELUPlugin::supportsFormatCombination判断该插件所支持的数据格式和类型，在这里我们告诉引擎，本插件可以支持什么类型的推理，可以支持多种，例如fp32、fp16、int8等等。
- 调用MySELUPlugin::getOutputDimensions获取该层的输出维度是多少。
- 调用MySELUPlugin::enqueue进行性能测试（不是一定会执行），如果支持多种，则会在多种里面进行实际测试，选择一个性能最好的配置。
- 调用MySELUPlugin::configurePlugin配置插件格式，告诉你目前这个层所采用的数据格式和类型。
- 调用MySELUPlugin::serialize将该层的参数序列化储存为trtmodel文件

推理阶段
- 通过MySELUPluginCreator::deserializePlugin反序列化插件参数进行创建。
- 期间会调用MySELUPlugin::clone克隆插件。
- 调用MySELUPlugin::configurePlugin配置当前插件使用的数据类型和格式。
- 调用MySELUPlugin::enqueue进行推理。

（三）简单快捷的plugin

上面的方法是官方给的方法，编写起来比较繁琐，我们可以对插件进行封装，使得用起来更加简单。在导出onnx模型，symbolic函数返回时，g.op返回的永远是Plugin这个名字，然后name_s指定为自己注册的插件名称，info_s则传递为json字符串，那么复合属性就可以轻易得到支持。

那么我们就可以省略上述的很多文件，我们只需要一个.cu文件就可以了。

那么在.cu文件如何注册一个插件呢：

class MYSELU : public TRTPlugin {
public:
	SetupPlugin(MYSELU);

	virtual void config_finish() override{
		printf("\033[33minit MYSELU config: %s\033[0m\n", config_->info_.c_str());
		printf("weights count is %d\n", config_->weights_.size());
	}

	int enqueue(const std::vector<GTensor>& inputs, std::vector<GTensor>& outputs, const std::vector<GTensor>& weights, void* workspace, cudaStream_t stream) override{
		
		int n = inputs[0].count();
		const int nthreads = 512;
		int block_size = n < nthreads ? n : nthreads;
		int grid_size = (n + block_size - 1) / block_size;

		MYSELU_kernel_fp32 <<<grid_size, block_size, 0, stream>>> (inputs[0].ptr<float>(), outputs[0].ptr<float>(), n);
		return 0;
	}
};

RegisterPlugin(MYSELU);

那么在这么便捷的背后其实就是写了很多默认值，不一样的就只有enqueue函数，从一个插件的创建到执行，都是使用写好的类来进行。

然后对应的builtin_op_importers.cpp添加了Plugin的解析支持：

DEFINE_BUILTIN_OP_IMPORTER(Plugin)
{
    std::vector<nvinfer1::ITensor*> inputTensors;
    std::vector<onnx2trt::ShapedWeights> weights;
    for(int i = 0; i < inputs.size(); ++i){
        auto& item = inputs.at(i);
        if(item.is_tensor()){
            nvinfer1::ITensor* input = &convertToTensor(item, ctx);
            inputTensors.push_back(input);
        }else{
            weights.push_back(item.weights());
        }
    }

    OnnxAttrs attrs(node, ctx);
    auto name = attrs.get<std::string>("name", "");
    auto info = attrs.get<std::string>("info", "");

    // Create plugin from registry
    auto registry = getPluginRegistry();
    auto creator = registry->getPluginCreator(name.c_str(), "1", "");
    if(creator == nullptr){
        printf("%s plugin was not found in the plugin registry!", name.c_str());
        ASSERT(false, ErrorCode::kUNSUPPORTED_NODE);
    }
    
    nvinfer1::PluginFieldCollection pluginFieldCollection;
    pluginFieldCollection.nbFields = 0;

    ONNXPlugin::TRTPlugin* plugin = (ONNXPlugin::TRTPlugin*)creator->createPlugin(name.c_str(), &pluginFieldCollection);
    if(plugin == nullptr){
        LOG_ERROR(name << " plugin was not found in the plugin registry!");
        ASSERT(false, ErrorCode::kUNSUPPORTED_NODE);
    }

    std::vector<std::shared_ptr<ONNXPlugin::Weight>> weightTensors;
    for(int i = 0; i < weights.size(); ++i){
        auto& weight = weights[i];
        std::vector<int> dims(weight.shape.d, weight.shape.d + weight.shape.nbDims);
        std::shared_ptr<ONNXPlugin::Weight> dweight(new ONNXPlugin::Weight(dims, ONNXPlugin::DataType::Float32));
        
        if(weight.type != ::onnx::TensorProto::FLOAT){
            LOG_ERROR("unsupport weight type: " << weight.type);
        }
        
        memcpy(dweight->pdata_host_, weight.values, dweight->data_bytes_);
        weightTensors.push_back(dweight);
    }
    
    plugin->pluginInit(name, info, weightTensors);
    auto layer = ctx->network()->addPluginV2(inputTensors.data(), inputTensors.size(), *plugin);
    std::vector<TensorOrWeights> outputs;
    for( int i=0; i< layer->getNbOutputs(); ++i )
      outputs.push_back(layer->getOutput(i));
    return outputs;
}
} // namespace

在这个函数中我们可以发现他可以读取插件名称以及通过json读取插件信息，之后创建相应的插件，之后我们就需要新建一个onnxplugin.cpp，实现对IPluginV2DynamicExt的封装。实现需要的函数有如下：

config_finish[非必要]：配置完成函数
- 当插件配置完毕时调用，可以在其中拿到各种属性，例如info、weights等

new_config[非必要]：实例化一个配置对象
- 可以自定义LayerConfig类并返回，也可以直接使用LayerConfig类
- 这个函数最大的作用，是配置本插件支持的数据格式和类型。比如fp32和fp16的支持等
getOutputDimensions[非必要]，获取该插件输出的shape大小，默认取第一个输入的大小
- 对应于原始插件的getOutputDimensions函数
enqueue[必要]，插件推理过程
- 插件的实际推理过程，该函数可能在编译和推理阶段数次调用

之后就可以注册插件，使用RegisterPlugin宏，格式是RegisterPlugin(类名);

0x07 INT8量化

int8量化是利用int8乘法替换float32乘法实现性能加速的一种方法。

对于常规模型有y = kx + b，此时x、k、b都是float32, 对于kx的计算使用float32的乘法。

对于int8模型有：y = tofp32(toint8(k) * toint8(x)) + b，其中int8 * int8结果为int16。

因此int8模型解决的问题是如何将float32合理的转换为int8，使得精度损失最小。

量化步骤：

config->setFlag(nvinfer1::BuilderFlag::kINT8);

实现Int8EntropyCalibrator类并继承自IInt8EntropyCalibrator2。

// 配置int8标定数据读取工具
shared_ptr<Int8EntropyCalibrator> calib(new Int8EntropyCalibrator(
    {"kej.jpg"}, input_dims, preprocess
));
config->setInt8Calibrator(calib.get());

实例化Int8EntropyCalibrator并且设置到config.setInt8Calibrator。
Int8EntropyCalibrator的作用，是读取并预处理图像数据作为输入。
- 标定的原理，是通过输入标定图像I，使用参数WInt8推理得到输出结果PInt8，然后不断调整WInt8，使得输出PInt8与PFloat32越接近越好
- 因此标定时通常需要使用一些图像，正常发布时，一般使用100张图左右即可

常用的Calibrator有如下：

参考：s7310-8-bit-inference-with-tensorrt.pdf (gputechconf.com)

Int8EntropyCalibrator2 熵校准选择张量的比例因子来优化量化张量的信息论内容，通常会抑制分布中的异常值。这是当前推荐的熵校准器。默认情况下，校准发生在图层融合之前。推荐用于基于 CNN 的网络。

Iint8MinMaxCalibrator

该校准器使用激活分布的整个范围来确定比例因子。它似乎更适合NLP任务。默认情况下，校准发生在图层融合之前。推荐用于NVIDIA BERT等网络。

计算机中的float计算量是非常打的，而改成int8后，计算量相比可以提升数倍。

对于实际操作时，input[float32], w[int8], bias[float32], output[float32]，步骤可以如下：

input[int8] = to_int8(input[float32])
y[int16] = input[int8] * w[int8] # 此处乘法会由计算机转换为int16，保证精度
output[float32] = to_float32(y[int16]) + bias[float32]

整个过程的只是为了减少float32的乘法数量以实现提速，对于to_int8的过程，并不是直接的线性缩放，而是经过KL散度计算最合适的截断点（最大、最小值），进而进行缩放，使得权重的分布尽可能小的被改变。