[TensorRT] 对UNet进行推理加速

        在过去的两周的实习时间里，公司有使用TensorRT对UNet进行推理加速的需求，目前我这边已经取得了一些进展，所以在这里稍微总结一下。

1. UNet网络

        网络有很多对UNet算法的说明，因为该部分并不是本文的重点，就不详细进行介绍了。为了更好的使用TensorRT的API搭建UNet的网络模型，本文只简单介绍UNet的网络架构，理清Encoder和Decoder阶段之间的关系。

         在上图中不难发现，UNet网络主要由三部分组成：下采样部分（Encoder）、上采样部分（Decoder）以及跨层连接部分（Concat）。在下采样部分主要使用的是最大池化以及卷积操作。使用步长为2的最大池化操作特征图尺寸的压缩，使用卷积操作进行特征提取（这种搭配方式，在前些年的神经网络中随处可见）；在上采样部分则依次使用Upsample操作进行特征图尺寸的扩展；在跨层连接部分中，主要会使用Concat操作，将编解码部分所对应的特征图进行融合并使用卷积操作对融合的特征图进行特征提取。

2. 使用TensorRT的API搭建UNet网络

        我们可以把TensorRT看做是类似PyTorch和TensorFlow之类的一种框架，只不过TensorRT使用较为苦涩难懂的C++语言进行编程，苦涩有苦涩的优点，但是它快啊！并且还能够自动进行网络层的计算优化，实现准确并快速的推理。

        所以，为了能够使用UNet网络进行推理加速，我们首先需要先创建UNet网络啊！TensorRT使用INetworkDefinition这个类进行网络的搭建，我对TensorRT搭建网络的理解就是先给你提供一个钩子，并且每个网路层的输入输出两个钩子，而你要做的就是将所有网络层按照一定的逻辑串起来就行。

        目前，INetworkDefiniton类中已经将常见的、简单的网络层实现了（比如卷积层、池化层、激活函数层等），我们只需要使用就可以了。对于BatchNormalization层等比较复杂的或者YOLO检测层等自定义的网络层暂时还没有提供接口可以使用，但可以重写TensorRT中现有的网络层进行自定义或者使用Plugin的形式，将某些网路层的计算逻辑重写出来，并将其作为网络中的某一个层进行使用。

        下面会先介绍一下在UNet使用到的并且也有现成接口的网络层：卷积层、激活函数层、Resize层、Padding层以及Concatenation层，然后再介绍一下需要稍微处理的BatchNormalization网络层。假设我们已经创建了一个INetworkDefinition的一个指针对象，我们现在要做的就是串羊肉串，不是，是给构建网络层。

2.1 Convolution层

        在TensorRT的API中，我们可以是使用addConvolutionNd()函数添加可以处理多维数据的卷积层，具体的函数接口如下所示：

IConvolutionLayer* INetworkDefinition::addConvolutionNd	(ITensor & input,
int32_t 	nbOutputMaps,
Dims 	kernelSize,
Weights 	kernelWeights,
Weights 	biasWeights 
)		

       函数的形参总共有5个，大致可以分为三个部分。第一部分是和输入数据相关的参数（Input变量），第二部分就是和输出数据相关的参数，需要设定输出特征图的个数以及卷积核的尺寸，第三部分就是该卷积处理的权重以及偏置。

​​​​​​​​​​​​​​

      一般我们设置卷积层的时候，都需要制定卷积核的尺寸（kernel size）、步长（stride）以及padding的大小。而上述添加卷积层的API中并没有参数制定strde以及padding，它是以函数方法的形式进行显示指定的。下面以一个案例来进行详细的说明：

       【案例1】给network添加kernel size为3， stride为2，padding为1的卷积层。

// 添加卷积层（weightMap为参数容器）
IConvolutionLayer* conv1 = network->addConvolutionNd(input, middle, DimsHW{k,k}, 
weightMap[lname + ".weight"], 
weightMap[lname +  ".bias"]);

assert(conv1);
conv1->setStrideNd(DimsHW{2, 2});
conv1->setPaddingNd(DimsHW{1, 1});

2.2 激活函数层

      激活函数层就比较简单了，我们只需要给定输入数据，并指定激活函数的类型就可以了。

IActivationLayer* INetworkDefinition::addActivation	(ITensor & input, 
ActivationType type )	

      目前，TensorRT中提供了12中常见的激活函数类型，使用的时候就只需要指定使用的激活函数类型就可以了，具体的激活函数类型可以参考这里。

      【案例2】给在案例1中的conv1网络层后面添加一个ReLU激活函数层。

// 添加ReLU网络层
IActivationLayer* relu1 = network->addActivation(*conv1->getOutput(0), 
ctivationType::kRELU);
assert(relu1);

2.3 Resize层

      Resize网络层能够实现对输入数据的缩放。因此在上采样过程中，可以用于对输入图像进行扩展。

IResizeLayer* INetworkDefinition::addResize (ITensor& 	input)	

      这个函数的使用也非常的简单，只需要传递输入变量就可以。只不过后续需要制定缩放的比例以及缩放所采用形式（线性插值或者最近邻插值）。

    【案例3】向案例2的relu1网络层中添加一个resize网络层，使用线性插值的方法，使得输出数据的尺寸是输入数据尺寸的2倍。

IResizeLayer* resizeLayer = network->addResize(input);
assert(resizeLayer);
// 使用线性插值的模型
resizeLayer->setResizeMode(ResizeMode::kLINEAR);
resizeLayer->setAlignCorners(true);
// 需要制定channel，heigh和widht三个通道的缩放比例
float scales[] = {1.0, 2.0, 2.0};
resizeLayer->setScales(scales, 3);

2.4 Padding层

      Padding网络层就只需要指定数据前后的padding的个数就行（默认使用0进行填充）。

IPaddingLayer* nvinfer1::INetworkDefinition::addPaddingNd(	ITensor & input,
Dims 	prePadding,
Dims 	postPadding )	

        其中需要说明的就是，prePadding以及postPadding。在TensorRT中，无论输入数据是几维的，我们都需要将输入数据拉平成一维数据，然后再进行处理。显而易见，对于一维数据的pre和post无非就是指开头和结尾。因此，对于上述函数中的prePadding以及postPadding就是指在一维数据的开头以及结尾添加的padding尺寸。

        【案例4】给案例3中的resizeLayer添加一个padding，prePadding以及postPadding都是2。

IPaddingLayer* pad1 = network->addPadding(*resizeLayer->getOutput(0), 
DimsHW(2, 2), DimsHW(2,  2));

2.5 Concat层

        Concat网络层就是将多个输入数据叠加在一起作为一个数据进行输出。在UNet以及YOLO检测模型中都有着非常好的检测效果，其本质就是将低维特征与高维语义特征进行融合，丰富高维特征的信息。

IConcatenationLayer* INetworkDefinition::addConcatenation(ITensor *const * 	inputs,
int32_t nbInputs )	

        该函数的输入是一个指针数组以及该指针数组中元素的个数。

        【案例5】 对tensor1和tensor2两个变量进行concate操作。

ITensor* tensors[] = {&tensor1, tensor2};
assert(tensors);
IConcatenationLayer* concat1 = network->addConcatenation(tensors, 2);

2.6 BatchNormalization层

        在TensorRT中并没有BN层，但是有一个IScaleLayer网络层，它能够对输入数据进行如下的处理：

        因此，我们可以对BN操作进行变化，从而让其符合IScaleLayer的格式。

        因此，IScalueLayer的参数和BN中的参数一一对应就为：

        在代码中所需要进行的操作就是得到BN中的四个参数，然后对其进行处理得到IScaleLayer中所需要的三个参数。

IScaleLayer* addBatchNorm2d(INetworkDefinition *network, 
std::map& weightMap, 
ITensor& input, std::string lname, float eps) {
    float *gamma = (float*)weightMap[lname + ".weight"].values;
    float *beta = (float*)weightMap[lname + ".bias"].values;
    float *mean = (float*)weightMap[lname + ".running_mean"].values;
    float *var = (float*)weightMap[lname + ".running_var"].values;
    int len = weightMap[lname + ".running_var"].count;
    // std::cout << "len " << len << std::endl;
    write(lname + ".weight");
    write(lname + ".bias");
    write(lname + ".running_mean");
    write(lname + ".running_var");
    float *scval = reinterpret_cast(malloc(sizeof(float) * len));
    for (int i = 0; i < len; i++) {
        scval[i] = gamma[i] / sqrt(var[i] + eps);
    }
    Weights scale{DataType::kFLOAT, scval, len};
    
    float *shval = reinterpret_cast(malloc(sizeof(float) * len));
    for (int i = 0; i < len; i++) {
        shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps);
    }
    Weights shift{DataType::kFLOAT, shval, len};

    float *pval = reinterpret_cast(malloc(sizeof(float) * len));
    for (int i = 0; i < len; i++) {
        pval[i] = 1.0;
    }
    Weights power{DataType::kFLOAT, pval, len};

    weightMap[lname + ".scale"] = scale;
    weightMap[lname + ".shift"] = shift;
    weightMap[lname + ".power"] = power;
    IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, 
scale, power);
    assert(scale_1);
    return scale_1;
}

      UNet网络搭建部分等之后再贴出来。

3 序列化保存模型

        在构建UNet网络之后，我们需要整个推理模型序列化为二进制数据保存到硬盘上，下次就可以直接加载这个模型进行推理就可以了，就不需要在从头开始重新构建模型了。

        序列化保存模型主要有以下七个步骤：

3.1 logger -> builder

        我们首先需要得到ILogger的一个对象，配置序列化模型过程中所需要的日志信息。

#include #include "NvInfer.h"

ILogger gLogger;
IBuilder* builder = createInferBuilder(gLogger);

3.2 builder -> config

IBuilderConfig* config = builder->createBuilderConfig();

3.3 builder -> network

INetworkDefinition* network = builder->createNetworkV2(0U);
assert(network);

3.4 build -> engine

builder->setMaxBatchSize(batch_size);
config->setMaxWorkspaceSize(1<<20);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

3.5 engine进行序列化

ICudaEngine* engine = createCudaEngine(builder, batch_size);
IHostMemory* modelStream = engine->serialize();
assert(modelStream);
// 将modeStream保存到文件中
std::ofstream f(engineFile);
assert(f.is_open() && "falid to open engineFile!");
f.write(reinterpret_cast(modelStream->data()), modelStream->size());
f.close();

4. 初始化模型并进行推理

        在第三步中，我们已经将engine模型进行序列化并保存到硬盘中了。在本小节中，我们需要对序列化的engine模型进行初始化（反序列化），然后将其加载到GPU上进行推理加速。

4.1 初始化模型

std::ifstream f(engineFile);
assert(f.is_open() && "fail to open engineFile ");
f.seekg(0, std::ios::end);
engineSize = f.tellg();
f.seekg(0, std::ios::beg);
assert(engineSize > 0 && "engine is empty!");
char* engineStream =  new char[engineSize];
f.read(engineStream, engineSize);
f.close();
assert(engineStream && "engineStream is nullptr!");
// 反序列化模型
char* engineStream = deserializeModel(engineFile, engineSize);

4.2 推理模型

4.2.1 logger -> runtime

IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime);

4.2.2 runtime -> engine

ICudaEngine* engine = runtime->deserializeCudaEngine(engineStream, 
engineSize, 
nullptr);
assert(engine);

4.2.3 engine -> context

IExecutionContext* context = engine->createExecutionContext();
assert(context);

4.2.4 inference

// 进行推理
void inference(IExecutionContext& context, float* input, float* output, const int batchsize){
    // 为输入输出分配显存
    std::cout << "inferencing... " << std::endl;
    const ICudaEngine& engine = context.getEngine();
    std::cout << "[ok] get engine!" << std::endl;
    // std::cout << engine.getName() << std::endl;
    int nBindings = engine.getNbBindings();
    assert(nBindings == 2);
    // 输入输出显存数组的指针
    void* buffers[2];
    const int inputIdx = engine.getBindingIndex(INPUT_BLOB_NAME);
    const int outputIdx = engine.getBindingIndex(OUTPUT_BLOB_NAME);
    // 为输入输出数组分配显存
    int memSize = batchsize * INPUT_CHANNEL * INPUT_HEIGHT * INPUT_WIDTH * sizeof(float);
    //int opmem = batchsize * 256*64*64 * sizeof(float);
    assert(memSize > 0);
    std::cout << "begine allocate mem ... " << std::endl;
    CHECK(cudaMalloc(&buffers[inputIdx], memSize));
    CHECK(cudaMalloc(&buffers[outputIdx], memSize));
    //CHECK(cudaMalloc(&buffers[outputIdx], opmem));
    std::cout << "[ok] finish cudaMalloc!" << std::endl;
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));

    // 将数据从host复制到device上
    std::cout << "copy memory from host to device ..." << std::endl;
    CHECK(cudaMemcpyAsync(buffers[inputIdx], input, memSize, cudaMemcpyKind::cudaMemcpyHostToDevice, stream));
    std::cout << "context enqueue ..." << std::endl;
    context.enqueue(batchsize, buffers, stream, nullptr);
    // 将计算结果从device复制到host上
    std::cout << "copy results from device to host ..." << std::endl;
    CHECK(cudaMemcpyAsync(output, buffers[outputIdx], memSize, cudaMemcpyKind::cudaMemcpyDeviceToHost, stream));

    //CHECK(cudaMemcpyAsync(output, buffers[outputIdx], memSize, cudaMemcpyKind::cudaMemcpyDeviceToHost, stream));
    // 同步，直到所有的stream都运行完
    CHECK(cudaStreamSynchronize(stream));
  
    //释放stream和显存
    std::cout << "free stream and memory.." <