【模型加速】CUDA-Pointpillars项目解读(2)
用于TensorRT的ONNX模型
接上篇【模型加速】CUDA-Pointpillars项目解读(1)，预处理之后输出pillar points bev特征图以及pillar coords。其size分别为(MAX_VOXELS,32,10)和(MAX_VOXELS,4)。原始的pointpillars网络结构中，预处理之后的pillar points bev要经过以下一系列操作(PointNet -> MaxPooling->)后送入Scatter层生成伪图像。
在cuda-pointpillars项目中作者对原有的网络结构进行了修改。其中ScatterBEV是单独开发了一个TensorRT的plugin，它顺带并入了ReduceMax的功能。所以造成了新老两种网络结构上的差异。
ScatterBEV Plugin核心部分的代码如下。包含3个输入和1个输出。输出即为我们要生成的伪图像。主要包含两大操作：ReduceMax以及ScatterBEV，分别以cuda核函数实现。
int ScatterBevPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace,
cudaStream_t stream) noexcept
{
cutelog("wow I run to here now");
unsigned int batch = 1;
unsigned int featureNum = featureNum_;
unsigned int featureY = feature_y_size_;
unsigned int featureX = feature_x_size_;
const float *in = (const float *)inputs[0];
const float *coords_data = (const float *)(inputs[1]);
const unsigned int *params_data = (const unsigned int *)(inputs[2]);
float *spatial_feature_data = (float *)(outputs[0]);
unsigned int count = inputDesc[0].dims.d[0];
cacheBEV_ = workspace;
const float *pillar_features_data = (const float *)(cacheBEV_);
//cudaMemcpyAsync(paramsPtr, params_data, 5*sizeof(int), cudaMemcpyDefault, stream);
checkCudaErrors(cudaMemsetAsync(spatial_feature_data, 0, batch*featureNum*featureY*featureX * sizeof(float), stream));
checkCudaErrors(reduceMax_kernel_launcher((const float*)in, (float*)pillar_features_data, count, stream));
checkCudaErrors(scatterBEV_kernel_launcher(pillar_features_data, coords_data, params_data, featureX, featureY, spatial_feature_data, stream));
return 0;
}最终,TensorRT推理部分的整体结构可以用作者给出的这张图来描叙。
为什么要把ReduceMax合并到ScatterBev中呢?因为该操作在TensorRT中相当耗时，我做了个实验，如果不做合并你再来看各层耗时。
load file: ../../data/000009.bin
find points num: 70148
find pillar_num: 22256
640 + (Unnamed Layer* 1) [Shuffle] 0.123ms
MatMul_245 input reformatter 0 2.624ms
MatMul_245 5.829ms
Transpose_246 + (Unnamed Layer* 4) [Shuffle] 14.978ms
BatchNormalization_247 7.158ms
(Unnamed Layer* 6) [Shuffle] + Transpose_248 13.937ms
Relu_249 7.091ms
ReduceMax_250 51.019ms
Squeeze_251 input reformatter 0 0.598ms
onnx_graphsurgeon_node_0 1.228ms
Conv_334 + Relu_335 input reformatter 0 1.110ms
Conv_334 + Relu_335 1.078ms
Conv_336 + Relu_337 0.866ms
Conv_338 + Relu_339 0.855ms
Conv_340 + Relu_341 0.849ms
Conv_360 + Relu_361 0.418ms
ConvTranspose_342 0.879ms
Conv_362 + Relu_363 0.693ms
Conv_364 + Relu_365 0.685ms
Conv_366 + Relu_367 0.690ms
Conv_368 + Relu_369 0.689ms
Conv_370 + Relu_371 0.693ms
Conv_390 + Relu_391 0.410ms
ConvTranspose_372 0.900ms
Conv_392 + Relu_393 0.739ms
Conv_394 + Relu_395 0.762ms
Conv_396 + Relu_397 0.761ms
Conv_398 + Relu_399 0.751ms
Conv_400 + Relu_401 0.736ms
ConvTranspose_402 1.279ms
BatchNormalization_343 + Relu_344 0.590ms
BatchNormalization_373 + Relu_374 0.588ms
BatchNormalization_403 + Relu_404 0.588ms
Conv_410 || Conv_407 || Conv_406 3.374ms
Conv_410 || Conv_407 || Conv_406 output reformatter 0 1.230ms
Transpose_411 0.417ms
Transpose_409 1.087ms
Transpose_408 1.092ms
其中MaxPooling最为耗时，这是在我的网络模型中测试的，和官方给出的稍有不同，但也足够说明问题。