最终用TensorRT重构centernet,由于python代码太过于复杂,单前处理就花了很长时间。现在前处理基本上没问题,网络已经可以在tensorrt上进行forward,但是真正复杂的地方才刚开始。

具体来说,centernet的后处理分为两个部分:

  1. 对 1x80x128x128这个feature进行3x3的maxpool,其实就等于是我们传统的进行nms一样;
  2. 从1x80x128x128里面根据热力图峰值点,找到我们需要的前100个目标,这里面有一个很重要的操作叫做TopK,这个topk就是将张量的最大的几个值返回,同时返回对应的index。然后根据index找到对应的中心点和偏移,对应的最终框也就找到了。

比较棘手的问题也是两个:

  1. 最后的这一步maxpool和topk为什么不能放到模型里面去?而是需要后处理来算?这机加大了复杂度,同时有很浪费时间。

思考如下:

这里的maxpool其实事可以放到模型里面去的,但是那样的化, 你网络出来的东西就比较多了,要出来1x80x100 同时还要对应的index。

这些张量的操作放到C++里面怎么都不太好做。