tensorRT是怎么构建和编译一个模型的

一、tensorrt基础概念

TensorRT的核心在于对模型算子的优化（合并算子、利用GPU特性选择特定核函数等多种策略），通过tensortRT，能够在Nvidia系列GPU上获得最好的性能。因此tensorRT的模型需要在目标GPU上实际运行的方式选择最优算法和配置（不跑跑，我哪知道怎么才最快，因为显卡是有高低档之分的，高档和低档最大的区别就是里面配置的加速器还有架构是非常不一样的，如果我们用tensorRT进行编译的话，那么它会根据你的显卡所开放的一些性能，充分得利用这些性能去使得速度尽可能快），也因此tensorRT生成的模型只能在特定条件下运行（编译的trt版本，cuda版本，编译时的GPU型号）

二、tensorrt的工作流程如下图

1.首先定义网络（这里的定义网络是指我现在有个模型我要告诉tensorrt，我通过什么方式去告诉它，比如我的权重是多少我怎么去让tensorrt知道）
2.给builder传入要优化的参数（network和config，config表示创建的构建配置即指定TensorRT应该如何优化模型，network为创建的网络定义），builder的过程就是指对模型算子的优化的过程。
3.通过builder生成engine,用于模型保存、推理等
4.engine可以通过序列化和逆序列化转化模型数据类型（序列化就是转化为二进制byte文件，加快传输速率），再进一步推动模型由输入张量到输出张量的推理。

Code Structure
1.定义 builder, config 和network，其中builder表示所创建的构建器，config表示创建的构建配置（指定TensorRT应该如何优化模型），network为创建的网络定义。
2.编译生成engine模型文件
3.序列化engine模型文件并存储

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

// cuda include，包含了cuda的头文件是因为tensorrt通常和cuda伴随出现，因为要做一些内存分配之类的东西
#include <cuda_runtime.h>

// system include
#include <stdio.h>

//TRTLogger继承了nvinfer1::ILogger这个类，日志类，主要负责记录tensorrt编译过程中出现的任何消息
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(){
    // 本代码主要实现一个最简单的神经网络 figure/simple_fully_connected_net.png 
     
    TRTLogger logger; // logger是必要的，用来捕捉warning和info等

    // ----------------------------- 1. 定义 builder, config 和network -----------------------------
    // 这是基本需要的组件
    //创建builder这个实例，形象的理解是你需要一个builder去build这个网络，网络自身有结构，这个结构可以有不同的配置
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
    // builder是指针，所以用->
    // 创建一个构建配置即一个config的实例，这个config主要解决的是我们要配置的batchsize是多少，是fp16还是int8等等各种配置信息都通过config对象来实现，config来指定TensorRT应该如何优化模型，tensorRT生成的模型只能在特定配置下运行
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
    // 创建网络定义，其中createNetworkV2(1)表示采用显性batch size，新版tensorRT(>=7.0)时，不建议采用0非显性batch size
    //所谓 显式 和 隐式 的差别就在于 Batch 这一维，即 显式 ==> NCHW，隐式 ==> CHW
    // 因此贯穿以后，请都采用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};

    //输入指定数据的名称（这里输入节点的名称就是image）、数据类型和完整维度，将输入层添加到网络
    //这里输入image必须为4个维度，所以我们定义为(1,3,1,1)
    //addInput几次就有几个输入
    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进行了解引用
    //添加激活层，它的输入就是layer1的输出，类型是kSIGMOID
    auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); // 注意更严谨的写法是*(layer1->getOutput(0)) 即对getOutput返回的指针进行解引用
    
    // 将我们需要的prob标记为输出，如果有多个输出要标记的话，就可以调用多次markOutput，把tensor（markOutput的输入参数）给塞进去就可以作为模型的输出
    //markOutput几次就有几次输出，markOutput是表示该模型的输出节点，为了节省内存，中间结果是不会被储存的即只有网络的输入和输出才储存
    network->markOutput(*prob->getOutput(0));
    //到此为止，network里面已经存好了模型的数据流转的结构和过程以及对应的参数
    //这时候我们可以开始配置整个编译的过程，配置MaxWorkspaceSize（指工作空间的大小，比如某些layer需要额外的存储时，它不会自己分配空间，而是为了内存复用直接找tensorRT要workspace空间。指的是这个意思）为256Mib，配置MaxBatchSize。1 << 28表示2的28次方
    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方法
    //通过buildEngineWithConfig方法把network和config塞进去编译得到engine！
    //编译的过程就是它去合并算子，去查找特定核函数等等，做各种优化等，做成功了那么engine就有了，如果失败了engine就是空指针
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if(engine == nullptr){
        printf("Build engine failed.\n");
        return -1;
    }

    // ----------------------------- 4. 序列化模型文件并存储 -----------------------------
    // 编译结束我们拿到engine这个实例后，那么下一步就应该要把它存储为文件
    // 将模型序列化，并储存为文件
    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;
}

输入，模型结构和输出的基本信息如下图所示：

Refer：官方文档参考部分 C++ API