TensorRT7 Onnx模型动态多batch问题解决

升级到Tensort7, ONNX模型创建网络和预测都需要用V2版本。需要动态的使用batch_size需要做一些额外的操作。

创建模型

    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    INetworkDefinition *network = builder->createNetworkV2(explicitBatch); // 创建TensorRT网络

预测模型

context_->enqueueV2(buffer_queue_.data(), stream_, nullptr);

这个遇到几个问题，用v1版本的创建和预测onnx模型会报必须指定explicitBatch。也就是说Tensort7解析解析了4维的第一维度。

换成V2的创建方式可以解决，但是可看到enqueueV2没有了batchsize这个参数。假设我们的onnx模型是10个batch,当只用到1个batch，我们填充数据用1个batch,预测的时候也按10个batch去infer。这样预测时间也是10个，达不到TensoRT7以前可以设定最大bathsize，实际按照小于max_bathsize的个数的运行。

为了解决这个问题，我们需要用到Tensor引入的动态Reshape的特性,这里我们解决第一维bathsize的问题。

首先onnx模型导出的时候第一维度需要用一个占位符号，而不是具体维度。下面是pytorch导出onnx模型代码片段

    input_names = ['input']
    output_names = ['output']
    #model  input_tensor  根据实际情况设置
    torch.onnx.export(model, input_tensor, 'example.onnx',
    export_params=True, verbose=False, input_names = input_names,
    output_names=output_names, opset_version=10,dynamic_axes={"input":{0:"batch_size"},
                                    "output":{0:"batch_size"}})

创建的时候需要创建优化配置类对象。根据输入尺寸设置为动态

    IBuilderConfig *builder_config = builder->createBuilderConfig();

    IOptimizationProfile *profile = builder->createOptimizationProfile();

    ITensor *input = network->getInput(0);

    Dims dims = input->getDimensions();

    profile->setDimensions(input->getName(),
                           OptProfileSelector::kMIN, Dims4{1, dims.d[1], dims.d[2], dims.d[3]});
    //max_batch_size_根据自己模型实际情况设置
    profile->setDimensions(input->getName(),
                           OptProfileSelector::kOPT, Dims4{max_batch_size_, dims.d[1], dims.d[2], dims.d[3]});
    profile->setDimensions(input->getName(),
                           OptProfileSelector::kMAX, Dims4{max_batch_size_, dims.d[1], dims.d[2], dims.d[3]});
    builder_config->addOptimizationProfile(profile);
    //创建cuda_engine 用带withconfig去创建
    ICudaEngine *engine = builder->buildEngineWithConfig(*network, *builder_config);

预测之前需要绑定维度

    //对应参数是batchsize  channel height width   这里只有batchsize是浮动的，其他三个就是网络的输出尺寸
    Dims4 input_dims{batch_size_, input_tensor_.channel(), input_tensor_.height(), input_tensor_.width()};
    context_->setBindingDimensions(0, input_dims);
    context_->enqueueV2(buffer_queue_.data(), stream_, nullptr);

注意初始化内部变量时候inputdim.d[0]由于是占位符号，值变成-1,这里用设置最大max_batch_size。输出维度信息也一样。

Dims inputdim = engine_->getBindingDimensions(b); // C*H*W
input_shape_.Reshape(max_batch_size_, inputdim.d[1], inputdim.d[2], inputdim.d[3]); // [batch_size, C, H, W]