4.4.tensorRT基础(1)-模型推理时动态shape的具体实现要点

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。

本次课程学习 tensorRT 基础-模型推理时动态 shape 的具体实现要点

课程大纲可看下面的思维导图

1. 动态shape

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

对于全卷积网络其实是有这么一个需求的，推理时输入 shape 可以动态改变的，不一定要限制死

动态 shape 案例代码如下：

// tensorRT include
#include 
#include 

// cuda include
#include 

// system include
#include 
#include 

#include  
#include  // 后面要用到ios这个库
#include 

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.网络结构：

首先是网络结构上的差异，上节课我们使用的是 Linear 层，这次我们使用 conv 层来替代，同时 activation 也修改成了 relu，通过 addConvolution 来添加卷积层

2.模型构建：

在模型构建阶段，动态 shape 模型的输入 shape 定义为 nvinfer1::Dims4(-1, num_input, -1, -1)，其中的 -1 表示该维度是动态的，即在运行时可以接受任何长度。

此外，动态 shape 模型还需要设置 Optimization Profile，即 profile 对象。这个对象定义了模型输入可能的最小、最大和最优 shape，在代码中 profile->setDimensions 方法用来设置

3.模型推理

在推理阶段，动态 shape 只需要在每次推理前设置输入数据的 shape 即可，使用 execution_context->setBindingDimensions 方法

运行结果如下所示：

图1-1 动态shape案例TRT推理结果

图1-2 动态shape案例Pytorch推理结果

可以看到 TRT 的输出和 Pytorch 一致，说明整个动态 shape 推理过程没问题

关于代码的重点提炼：

1. OptimizationProfile 是一个优化配置文件，它就是用来指定输入的 shape 可以变换的范围的，不要被优化两个字蒙蔽了双眼

2. 如果 onnx 模型的输入某个维度是 -1，则表示该维度是动态的，否则表示该维度是明确的，明确维度的 minDims，optDims，maxDims 一定是一样的。

关于动态 shape 的知识点有：(from 杜老师)

1. 构建网络时：

• 1.1. 必须在模型定义时，输入维度给定为 -1，否则该维度不会动态。注意两点：
• 1.1.1 若 onnx 文件，则 onnx 文件打开后应该看到为动态或者 -1
• 1.1.2 如果你的模型中存在 reshape 类型，那么 reshape 的参数必须随动态计算。而大部分时候都是问题，除非你是全卷积模型，否则大部分时候只需要为 batch_size 维度设置为动态，其它维度尽量避免设置动态、
• 1.2. 配置 profile：
• 1.2.1 create：builder->createOptimizationProfile()
• 1.2.2 set：setDimension() 设置 kMIN，kOPT，kMAX 的一系列输入尺寸范围
• 1.2.3 add：config->addOptimizationProfile(profile); 添加 profile 到网络配置中

2. 推理阶段时：

• 2.1. 关于 profile 索引

• 2.2 在运行时，向 engine 请求绑定维度会返回用于构建网络的相同维度。这意味着，得到的还是动态的维度 [-1, in_channel, -1, -1]：

1.multiple-optimization-profiles.jpg

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

1.multiple-optimization-profiles.jpg

3. 检查正确性

• 我们通常可以利用 pytorch 来校验是否发生了错误

2. 补充知识

我们绝大部分时候只考虑 batch 维度的动态，并不太关注宽高动态，关于静态 batch 和动态 batch 有以下几点说明：

静态 batch

• 导出的 onnx 指定所有维度均为明确的数字，是静态 shape 模型
• 在推理的时候，它永远是同样的 batch 推理，即使你目前只有一个图推理，它也需要 n 和 batch 的耗时
• 适用于大部分场景，整个代码逻辑非常简单

动态 batch

• 导出的时候指定特定维度为 dynamic，也就是不确定状态
• 模型推理时才决定所需推理的 batch 大小，耗时最优，但 onnx 复杂度提高了
• 适用于如 server 这种有大量不均匀的请求时的场景

更多细节请查看 https://www.bilibili.com/video/BV15Y41167B5/

总结

本次课程我们学习了动态 shape 的相关知识以及在 TRT 中实现的要点，在代码中我们主要通过优化配置文件 OptimizationProfile 来指定动态 shape 的。值得注意的是我们绝大部分情况下只会考虑 batch 维度的动态，静态 batch 推理的时候，永远是同样的 batch 推理，需要 batch 的耗时，适用于绝大部分场景；而动态 shape 只有在模型推理时才决定所需推理的 batch 大小，耗时最优，适用于服务器有大量不均匀的请求场景。