Pytorch转onnx到TensorRT部署流程

生产环境：ubuntu18.04

                  cuda 11.0

                  cudnn 8.0.5

 

一、注意事项

 1.1 bug记录

 1.1.1  Pytorch中的F.interpolate()函数

F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=True)

Pytorch的interpolate层可以转换成onnx，但从onnx转换成trt的时候，就会报错

ONNX IR version:  0.0.6
Opset version:    11
Producer name:    pytorch
Producer version: 1.7
Domain:           
Model version:    0
Doc string:       
----------------------------------------------------------------
Parsing model
While parsing node number 2 [Resize -> "output"]:
ERROR: /home/ego/GPU/onnx-tensorrt/builtin_op_importers.cpp:2616 In function importResize:
[8] Assertion failed: (transformationMode == "asymmetric" || transformationMode == "pytorch_half_pixel" || transformationMode == "half_pixel") && "TensorRT only supports half pixel, pytorch half_pixel, and asymmetric tranformation mode for linear resizes when scales are provided!"

这时就需要将scale_factor转换成size的类型，如下

F.interpolate(x, size=[int(2 * x.shape[2]), int(2 * x.shape[3])], mode=‘bilinear‘,align_corners=True)

转换生成的onnx模型，可以使用onnx-simplifier优化

python3 -m onnxsim input_onnx_model output_onnx_model

1.1.2 pytorch中的gather函数

torch中的gather函数，在onnx中以gatherElement的形式存在，在onnx中可以正常运行，但是在转成TensorRT时，TRT并不支持这种操作

GatherElement函数说明：和torch.gather类似

  out[i][j][k] = input[index[i][j][k]][j][k] if axis = 0,
  out[i][j][k] = input[i][index[i][j][k]][k] if axis = 1,
  out[i][j][k] = input[i][j][index[i][j][k]] if axis = 2,

example 1
data = [
      [1, 2],
      [3, 4],
  ]
  indices = [
      [0, 0],
      [1, 0],
  ]
  axis = 1
  output = [
      [
        [1, 1],
        [4, 3],
      ],
  ]

example 2
data = [
      [1, 2, 3],
      [4, 5, 6],
      [7, 8, 9],
  ]
  indices = [
      [1, 2, 0],
      [2, 0, 0],
  ]
  axis = 0
  output = [
      [
        [4, 8, 3],
        [7, 2, 3],
      ],
  ]

 

说明：当axis为1时，按行操作，对于【1，2】，选取[0,0]，那么就是[1,1]

对于【3，4】，选取【1，0】那么就是【4，3】

当axis=0时，按列操作，对于第一列，索引对应的就是【4，8，3】，第二列【7，2，3】

1.2 函数说明

1.2.1关于onnx格式的解析

二、TensorRT的量化

2.1 量化原理

2.1.1 非饱和量化

一般对于网络的权重参数使用非饱和量化，也就是直接用最大值/128,因为其分布比较对于激活值则使用饱和量化

2.1.2 饱和量化

2.1.3 模型选择

2.2 工作流程

宏观处理流程如下，首先准备一个校准数据集，然后对每一层：

• 收集激活值的直方图；
• 基于不同的阀址产生不同的量化分布；
• 然后计算每个分布与原分布的相对熵，然后选择熵最少的一个，也就是跟原分布最像的一个。

此时阀值就选出来啦，对应的scale值也就出来了。

而其中最关键的就是校准算法部分了：

calibration：基于实验的迭代搜索阀值。

校准是其核心部分，应用程序提供一个样本数据集（最好是验证集的子集），称为“校准数据集”，它用来做所谓的校准。

在校准数据集上运行FP32推理。收集激活的直方图，并生成一组具有不同阈值的8位表示法，并选择具有最少kl散度的表示；kl-散度是在参考分布（即FP32激活）和量化分布之间（即8位量化激活）之间。

校准数据集可以按照如下的方式制作

ls -R /home/datalab/work/datasets/test_7pilang/*.jpg > file.txt

首先看上图的原理，就是把大范围的一个值给缩小到一个小范围的值(注意是等比例的缩小)。

这里的2048bins指的是正范围还是负范围呢？因为我看到后面都是量化的128bins里面去的，也就是说只管了int8(256)的一半！具体细节下篇文章中的代码实现部分详细分析；

这里看他的意思就是输入为[0, 2048] bins，然后想办法把这么大的分布给找到一个合理的阀值T然后把阀值内的bins映射到int8的128个bins里面来，最终而且信息熵损失是最少的。

流程如下所示：

1 首先不断地截断参考样本P，长度从128开始到2048，为什么从128开始呢？因为截断的长度为128的话，那么我们直接一一对应就好了，完全不用衰减因子了；

这里bin的宽度就是FP32激活函数值的最大max/2048,然后为了找截断值T，就让长度从128~2048迭代；

2 统计截断区外的数值并将其算作bin[i-1]的count,方便计算概率分布

3 计算样本P的概率分布

4 创建样本Q，其元素的值为截断样本P的int8量化值

是按照下图计算的，比如样本P有256个bin，但样本Q只有128个bin

就将P样本 merge into 到128个bins，就是将相邻的两个数相加

5 扩展样本Q的分布到i个bin

也是如下图所示

6计算Q的概率分布并求他们的KL散度

不断循环，找到最小KL散度的数值(此时i=m），此时表示Q能极好地拟合P分布了

此时阈值就等于(m+0.5)*bin的长度;

三、TensorRT部署流程（基于ONNX）

注：主要参考TensorRT官方的SampleONNXmnist教程

3.1 Build阶段

3.1.1 Create the builder , the network and BuilderConfig

 auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
 const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);//left move 
 auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));// second create Network
 auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());

IBuilderConfig 可以控制网络推理在何种精度上(如FP16、INT8），也可以autotunig Tensor运行多少次内核来计算最快时间的参数，主要目的还是设置maximum workspace size

附maximum workspace size：限制某一层在network中能够使用的最大空间，一般设置为1<<20

3.1.2 Create parser and Using it to create the Onnx MNIST Network(Tensor RT)

auto parser= SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
auto constructed = constructNetwork(builder, network, config, parser);

constructNetwork函数如下所示：输入参数中的network指向输出的生成的onnx网络

//! \brief Uses a ONNX parser to create the Onnx MNIST Network and marks the
//!        output layers
//!
//! \param network Pointer to the network that will be populated with the Onnx MNIST network
//!
//! \param builder Pointer to the engine builder
//!
bool SampleOnnxMNIST::constructNetwork(SampleUniquePtr<nvinfer1::IBuilder>& builder,
    SampleUniquePtr<nvinfer1::INetworkDefinition>& network, SampleUniquePtr<nvinfer1::IBuilderConfig>& config,
    SampleUniquePtr<nvonnxparser::IParser>& parser)
{
    auto parsed = parser->parseFromFile(locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
        static_cast<int>(sample::gLogger.getReportableSeverity()));
    if (!parsed)
    {
        return false;
    }

    config->setMaxWorkspaceSize(16_MiB);
    if (mParams.fp16)
    {
        config->setFlag(BuilderFlag::kFP16);
    }
    if (mParams.int8)
    {
        config->setFlag(BuilderFlag::kINT8);
        samplesCommon::setAllTensorScales(network.get(), 127.0f, 127.0f);
    }

    samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);

    return true;
}

3.1.3 用设置的config参数来创建TensorRT网络

//When the engine is built, TensorRT makes copies of the weights
mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config), samplesCommon::InferDeleter());

3.2 Infer阶段

先申请缓存，然后设定输入，最后执行engine

    samplesCommon::BufferManager buffers(mEngine);
    auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine>createExecutionContext());
    // Read the input data into the managed buffers
    assert(mParams.inputTensorNames.size() == 1);
    if (!processInput(buffers))
    {
        return false;
    }

    // Memcpy from host input buffers to device input buffers
    buffers.copyInputToDevice();

    bool status = context->executeV2(buffers.getDeviceBindings().data());
    if (!status)
    {
        return false;
    }

    // Memcpy from device output buffers to host output buffers
    buffers.copyOutputToHost();
    // Verify results
    if (!verifyOutput(buffers))
    {
        return false;
    }

    return true;

缓存申请中有这样一个成员函数mEngine->getBindingVectorizedDim(),

Return the dimension index that the buffer is vectorized.

Specifically -1 is returned if scalars per vector is 1.

3.3 序列化模型及反序列化模型

为什么要序列化模型，因为build engine非常耗时

        IHostMemory *gieModelStream = engine->serialize();
        std::string serialize_str;
        std::ofstream serialize_output_stream;
        serialize_str.resize(gieModelStream->size());   
        memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size());
        serialize_output_stream.open("./serialize_engine_output.trt");
        serialize_output_stream<<serialize_str;
        serialize_output_stream.close();