升级到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]