【3D 目标检测模型部署】全链条打通！PointPillars 从模型到部署
目录
1. PointPillars 模型介绍
模型结构
Voxel-base 模型中使用到的 Voxelization
PillarFeatureNet 模块介绍
PoinPillars 在 MMDetection3D 中的实现
2. 快速部署 MMDetection3D 的 PointPillars 模型
常见的推理后端
模型中间件 ONNX
移出自定义算子
修改模块
总结
3D 目标检测是现今计算机视觉领域算法的一个重要分支，在许多领域都有着重要的应用，比如自动驾驶、智能机器人等。PointPillars 是 3D 目标检测算法中一个十分经典的模型，其较为简单的模型结构非常适合部署落地，使得其在实际场景中有着十分广泛的应用。那么今天我们就来介绍一个 PointPillars 模型以及从模型到部署的全过程。
MMDeploy 支持对 MMDetection3D 中 PointPillars 模型的一键转换部署，欢迎大家使用。
https://github.com/open-mmlab/mmdetection3dgithub.com/open-mmlab/mmdetection3d
https://github.com/open-mmlab/mmdeploygithub.com/open-mmlab/mmdeploy
1. PointPillars 模型介绍
近几年点云目标检测一直是计算机视觉领域中很火的任务，各种方法也是层出不穷，从基于点云的 PointNet、PointNet++，到使用体素方法的 VoxelNet。这些网络取得了不错的检测性能，但同时也存在一些问题。
PointNet 系列的 Point-based 模型直接对点云进行处理，可以减少位置信息的损失，但同时也带来了巨大的计算资源消耗，使其很难做到实时。VoxelNet 等 Voxel-base 的模型相较于 Point-base 的模型在推理速度上有所提升，但是由于模型中使用了三维卷积的 backbone，所以也仍然很难做到实时。
为了解决点云目标检测模型的速度问题，nuTonomy 公司于 2019 年提出了 PointPillars 模型，现如今 PointPillars 已经是最常用于点云目标检测任务的模型之一。相较于其他的模型，PointPillars 在推理速度方面有着明显的优势（遥遥领先），同时又能保持着不错的准确性，如下图所示。在配置为 Intel i7 CPU 和 1080ti GPU 上的推理速度为 62Hz，可以达到实时的效果，所以常常被应用在无人驾驶领域，是一个落地且应用广泛的 3D 快速目标检测网络，那么我们来了解一下 PointPillars 模型。
模型结构
首先我们来看一下 PointPillars 模型整体结构，如下图所示，PointPillars 的整个模型结构还是和 VoxelNet 类似。
整个网络结构可以分为三个部分：
- Pillar Feature Net：将输入的点云转换为稀疏的伪图像的特征形式。
- Backbone（2D CNN）：使用 2D 的 CNN 处理伪图像特征得到高维度的特征。
- Detection Head（SSD）：检测和回归 3D 边界框。
Voxel-base 模型中使用到的 Voxelization
我们先来看一下 voxel-base 的模型中常常使用到的 voxelization（体素化）。在实际使用过程中我们都希望我们的模型又快有准，所以为了可以权衡速度和精度，VoxelNet 提出了使用 voxelization（体素化）的方法来处理点云。点云是三维空间中的物体表示，因此一个自然的思路是将空间在长宽高三个方向划分格子，每个格子称为 voxel（体素），通过处理将其转换为 3 维数组的形式，再使用 3D 卷积和 2D 卷积的网络处理，如下图所示。
然而体素化也会带来一些问题，例如不可避免的会造成一些信息的丢失，对体素参数较为敏感，以及转换成 3 维数组后提取特征时通常需要用到 3 维卷积，但我们知道 3 维的卷积是一个相当耗时的操作，所以当我们设置体素化的粒度过大时会导致较多的信息丢失，但如果粒度过小又会导致计算时间几何增加。那么有什么办法在保证模型精度的同时，又能有高效的推理速度呢？
PointPillars 在 VoxelNet 中的 voxel 的基础上提出了一种改进版本的点云表征方法 pillar，可以将点云转换成伪图像的形式，进而通过 2D 卷积实现目标检测，相较于 VoxelNet 将点云转换成 voxel 形式然后使用相当耗时的 3 维卷积来处理特征，PointPillars 这种使用 2 维卷积的网络在推理速度上有很大的优势。那么什么是 pillar？原文中的描述是“ a pillar is a voxel with unlimited spatial extent in the z direction ”，其实很简单，将空间的 x，y 轴两个方向上划分格子，然后再将每个格子在 z 轴上拉伸，使其可以覆盖整个空间 z 轴，就可以得到一个 pillar，且空间中的每个点都可以划分到某个 pillar 中。
PillarFeatureNet 模块介绍
我们来看一下 PointPillars 中的 PillarFeatureNet 模块的具体操作，如下图所示。
- 首先将一个样本的点云空间划分成（在 X 轴方向上点云空间的范围 / pillar size，在 Y 轴方向上点云空间的范围 / pillar size）pillar 网格，样本中的点根据会被包含在各个 pillar 中，没有点的 pillar 则视为空 pillar。
- 假设样本中包含的非空 pillar 数量为 P，同时限制每个 pillar 中的点的最大数量为 N，如果一个 pillar 中点的数量不及 N，则用 0 补全，若超过 N，则从 pillar 内的点中采样出 N 个点来。并对 pillar 中的每个点进行编码，其中每个点的表示会包括点的坐标，反射强度，pillar 的几何中心，点与 pillar 几何中心的相对位置，将每个点的表示的长度记为 D。这样我们的一个点云样本就可以用一个（P，N，D）的张量来表示。
- 得到点云的 pillar 表示的张量后，我们对其进行处理提取特征，通过使用简化版的 PointNet 中的 SA 模块来处理每个 pillar。即先对每个 pillar 中的点使用多层 MLP 来使得每个点的维度从 D 变成 C，这样张量变成了（P，N，C），然后对每个 pillar 中的点使用 Max Pooling，得到每个 pillar 的特征向量，也使得张量中的 N 的维度消失，得到了（P，C）维度的特征图。
- 最后将（P，C）的特征根据 pillar 的位置展开成伪图像特征，将 P 展开为（H，W）。这样我们就获得了类似图像的（C，H，W）形式的特征表示。
PoinPillars 在 MMDetection3D 中的实现
MMDetection3D 中实现了 PointPillars 算法，以一个 kitti 数据形式的 PointPillars 模型作为例子，将模型拆分成了 6 个模块：voxel_layer，PillarFeatureNet，PointPillarsScatter，Backbone，FPN，Anchor3DHead。整体算法流程如下图所示：
2. 快速部署 MMDetection3D 的 PointPillars 模型
考虑到算法模型在实际使用时的对速度的需求和对不同设备的适配性，通常我们会将训练好的模型部署到某个推理后端上，那么我们先来简单介绍下目前比较常见的一些推理后端。
常见的推理后端
TensorRT 针对 NVIDIA 系列显卡具有其他框架都不具备的优势，如果运行在 NVIDIA 显卡上， TensorRT 一般是所有框架中推理最快的。一般的主流训练框架如 TensorFlow 和 Pytorch 都能将训练好的模型转换成 TensorRT 可运行的模型。当然了，TensorRT 的限制就是只能运行在 NVIDIA 显卡上，并且 kernel 未开源。
ONNXRuntime 是可以运行在多平台(Windows，Linux，Mac，Android，iOS)上的一款推理框架，它接受 ONNX 格式的模型输入，支持 GPU 和 CPU 的推理。唯一不足就是 ONNX 节点粒度较细，推理速度有时候比其他推理框架如 TensorRT 较低。
OpenVINO 是 Intel 家推出的针对 Intel 出品的 CPU 和 GPU 友好的一款推理框架，同时它可以对接不同训练框架如 TensorFlow，Pytorch，Caffe 等。不足之处可能是只支持 Intel 家的硬件产品。
NCNN，MNN 不同于以上三款推理框架，这两款框架都是针对手机端的部署。ncnn 是腾讯开源的，MNN 是阿里开源的。ncnn 的优势是开源较早，有非常稳定的社区，开源影响力也较高。mnn 开源略晚，但也是目前比较有影响力的手机端推理框架。
我们现在知道了一些推理后端，那么怎么把我们使用深度学习训练框架（如Pytorch，TensorFlow 等）训练出来的模型部署到这些推理后端上呢？这里我们就需要用到一些中间件（IR）来做为深度学习框架和推理后端之间的桥梁，例如 ONNX。
模型中间件 ONNX
Open Neural Network Exchange (ONNX) 是一个开放的生态系统，是一种深度模型的开放格式，它使开发人员能够随着项目的发展选择合适的工具，增强模型的交互性。 它定义了一个可扩展的计算图模型，内置运算符以及标准数据类型。ONNX 可以实现不同框架之间的互操作性并简化从研究到生产的路径，不同框架训练出来的模型都可以转换成 ONNX 模型进行存储以及后续的推理，目前 ONNX 在许多框架、工具和硬件中都得到了广泛的支持，是模型部署中最重要的中间表示之一。
目前 Pytorch 也提供了对 ONNX 的支持，可以参考 torch.onnx 文档。torch.onnx.export 是 PyTorch 自带的把训练好的 Pytorch 模型转换成 ONNX 模型的函数，这个函数的本质是使用跟踪（trace）的方法，通过跟踪一个具体的输入在模型中的运行过程，记录运行过程中的每个计算节点从而得到一个静态的计算图，进而生成一个 ONNX 模型文件。具体的 torch.onnx.export 如下：
def export(model, args, f, export_params=True, verbose=False, training=TrainingMode.EVAL,
input_names=None, output_names=None, aten=False, export_raw_ir=False,
operator_export_type=None, opset_version=None, _retain_param_name=True,
do_constant_folding=True, example_outputs=None, strip_doc_string=True,
dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None,
enable_onnx_checker=True, use_external_data_format=False): 我们来介绍一下一些可能需要用到的参数的具体意义：
model：我们要导出成 ONNX 模型的 Pytorch 模型。
args：模型的输入，也就是 torch.onnx.export 用来跟踪和运行模型得到静态计算图时使用的输入数据。
f：导出的 ONNX 文件名，如 "PointPillars.onnx"。
input_names, output_names：设置输入和输出张量的名称，ONNX 模型的每个输入和输出张量都有一个名字。例如我们后面导出 PointPillars 时会用到的输入输出名称：
input_names=['voxels', 'num_points', 'coors'],
output_names=['scores', 'bbox_preds', 'dir_scores'] dynamic_axes：指定输入输出张量的哪些维度是动态的。为了追求效率，ONNX 默认所有参与运算的张量都是静态的，即都是固定 shape 的输入输出。但是这和我们在实际使用时的一些需求是相违背的，所以我们可以使用 dynamic_axes 来指定输入输出的哪个维度是动态的，即可变大小的输入输出。由于点云数据的不规则性，无法确定输入点云的数量因此也无法确定 Pillar 的数量，所以我们在导出 ONNX 时需要注意要使用动态维度来指定哪些维度的大小是可变的，例如我们后面在导出 PointPillars 时用到的动态维度：
dynamic_axes={
'voxels': {
0: 'voxels_num',
},
'num_points': {
0: 'voxels_num',
},
'coors': {
0: 'voxels_num',
} 更多有关 Pytorch 到 ONNX 模型转换的知识可以参考我们的模型部署入门教程(三）：PyTorch 转 ONNX 详解。
好了现在我们已经了解了一些基础知识，那么来看如何部署一个 MMDetection3D 的 PointPIlars 模型。通常我们部署一个训练好的 Pytorch 模型会经过 Pytorch 模型-->ONNX 模型-->可以在推理后端上运行的模型 的转换过程。但是过程往往不会一帆风顺，我们一点一点的来解决问题。
移出自定义算子
首先在 PointPillars 的 ONNX 模型导出时我们会遇到一些问题，由于模型的 Voxel_layer 和 Anchor3DHead 的后处理中通常会使用自定义 Voxelize 算子和 3D NMS 算子来进行加速，如果想要将这些算子作为模型的一部分导出并在推理后端上推理，我们需要根据不同推理后端的要求添加相应的算子，这无疑是件十分麻烦的事情，因此我们可以选择将使用了 Voxelize 的 Voxel_layer 和使用了 3D NMS 的后处理从模型中拆分出来，从而使得导出的 ONNX 模型可以直接在不同的后端上运行推理，如下图所示。
为此，我们需要重写一些模型中的代码来将 Voxel_layer 和后处理从模型推理的流程中移出，我们需要修改一下 detector，示例代码如下（节选自 MMDeploy）：
# origin func
def simple_test(self, points, img_metas, imgs=None, rescale=False):
"""Test function without augmentation."""
x = self.extract_feat(points, img_metas)
outs = self.bbox_head(x)
bbox_list = self.bbox_head.get_bboxes( # 后处理，包含NMS算子
*outs, img_metas, rescale=rescale)
...
def extract_feat(self, points, img_metas=None):
"""Extract features from points."""
voxels, num_points, coors = self.voxelize(points) # 点云体素化，包含voxelize算子
voxel_features = self.voxel_encoder(voxels, num_points, coors)
batch_size = coors[-1, 0].item() + 1
...
# rewrite func
def simple_test(self,
voxels,
num_points,
coors,
img_metas=None,
imgs=None,
rescale=False):
x = self.extract_feat(voxels, num_points, coors, img_metas)
bbox_preds, scores, dir_scores = self.bbox_head(x) # 移除了后处理，直接输出head的结果
return bbox_preds, scores, dir_scores
def extract_feat(self, voxels, num_points, coors,
img_metas=None):
"""Extract features from points."""
voxel_features = self.voxel_encoder(voxels, num_points, coors) # 移除了voxelize，
batch_size = coors[-1, 0].item() + 1 # 把点云voxelize后的数据作为输入
... 这样我们可以在模型外调用函数来实现 Voxelize 和模型后处理。如下面的代码所示：
def voxelize(model_cfg: Union[str, mmcv.Config], points: torch.Tensor):
from mmcv.ops import Voxelization
model_cfg = load_config(model_cfg)[0]
voxel_layer = model_cfg.model['voxel_layer']
voxel_layer = Voxelization(**voxel_layer)
voxels, coors, num_points = [], [], []
for res in points:
res_voxels, res_coors, res_num_points = voxel_layer(res)
voxels.append(res_voxels)
coors.append(res_coors)
num_points.append(res_num_points)
voxels = torch.cat(voxels, dim=0)
num_points = torch.cat(num_points, dim=0)
coors_batch = []
for i, coor in enumerate(coors):
coor_pad = F.pad(coor, (1, 0), mode='constant', value=i)
coors_batch.append(coor_pad)
coors_batch = torch.cat(coors_batch, dim=0)
return voxels, num_points, coors_batch
def post_process(model_cfg: Union[str, mmcv.Config],
outs: torch.Tensor,
img_metas: Dict,
device: str,
rescale=False):
from mmdet3d.core import bbox3d2result
from mmdet3d.models.builder import build_head
model_cfg = load_config(model_cfg)[0]
head_cfg = dict(**model_cfg.model['bbox_head'])
head_cfg['train_cfg'] = None
head_cfg['test_cfg'] = model_cfg.model['test_cfg']
head = build_head(head_cfg)
cls_scores = [outs['scores'].to(device)]
bbox_preds = [outs['bbox_preds'].to(device)]
dir_scores = [outs['dir_scores'].to(device)]
bbox_list = head.get_bboxes(
cls_scores, bbox_preds, dir_scores, img_metas, rescale=False)
bbox_results = [
bbox3d2result(bboxes, scores, labels)
for bboxes, scores, labels in bbox_list
]
return bbox_results
好的，现在 PointPillars 模型中没有自定义算子了，我们可以试着来导出 ONNX 并转换到推理后端上来看一下推理结果，显然结果并不正确，如下图所示，我们来看看有哪些问题～
修改模块
我们一个一个模块的来解决问题，首先是 PillarFeatureNet 模块，我们单独导出它的 ONNX 模型来看一下，可以发现节点复杂繁多，如下图所示，而且该模块在推理后端上的推理结果和在 Pytorch 上的推理结果也不同。
通过逐步定位，我们发现问题出现在求 pillar 中的点和 pillar 的几何中心的相对位置上。
# origin func
def forward(self, features, num_points, coors):
...
if self._with_voxel_center:
if not self.legacy:
f_center = torch.zeros_like(features[:, :, :3])
f_center[:, :, 0] = features[:, :, 0] - (
coors[:, 3].to(dtype).unsqueeze(1) * self.vx +
self.x_offset)
f_center[:, :, 1] = features[:, :, 1] - (
coors[:, 2].to(dtype).unsqueeze(1) * self.vy +
self.y_offset)
f_center[:, :, 2] = features[:, :, 2] - (
coors[:, 1].to(dtype).unsqueeze(1) * self.vz +
self.z_offset)
else:
f_center = features[:, :, :3]
f_center[:, :, 0] = f_center[:, :, 0] - (
coors[:, 3].type_as(features).unsqueeze(1) * self.vx +
self.x_offset)
f_center[:, :, 1] = f_center[:, :, 1] - (
coors[:, 2].type_as(features).unsqueeze(1) * self.vy +
self.y_offset)
f_center[:, :, 2] = f_center[:, :, 2] - (
coors[:, 1].type_as(features).unsqueeze(1) * self.vz +
self.z_offset)
features_ls.append(f_center)
... 原先的函数使用了取下标的方式进行计算使得整个计算图显得十分复杂，同时这种赋值方式导致运算在 Pytorch 上表现为浅拷贝，而在一些推理后端上却表现为深拷贝，因此会造成结果差异，我们修改原先的代码，使用矩阵切片代替原先的操作，使得导出的模型在推理后端上的行为结果能和 Pytorch 一致，并简化了计算图：
# rewrite func
def forward(self, features, num_points, coors):
...
if self._with_voxel_center:
if not self.legacy:
f_center = features[..., :3] - (
coors[..., 1:] * torch.tensor([self.vz, self.vy, self.vx]).to(device)
+
torch.tensor([self.z_offset, self.y_offset, self.x_offset
]).to(device)).unsqueeze(1).flip(2)
else:
f_center = features[..., :3] - (
coors[..., 1:] * torch.tensor([self.vz, self.vy, self.vx]).to(device)
+
torch.tensor([self.z_offset, self.y_offset, self.x_offset
]).to(device)).unsqueeze(1).flip(2)
features_ls[0] = torch.cat((f_center, features[..., 3:]), dim=-1)
features_ls.append(f_center)
... 修改后我们发现计算图变得更加简洁，如下图所示，同时模块推理结果也和 Pytorch 的结果对齐。
好的，我们现在导出了一个“看起来”还可以的 ONNX 模型。我们试着将转换出来的 ONNX 模型部署到 TensorRT 上，发现报错：
[TRT] [I] No importer registered for op: NonZero. Attempting to import as plugin.
[TRT] [I] Searching for plugin: NonZero, plugin_version: 1, plugin_namespace:
[TRT] [E] 3: getPluginCreator could not find plugin: NonZero version: 1 原来是模型中有 Nonzero 算子，而有些后端没有支持 Nonzero，例如 TensorRT，我们对代码进行定位，发现问题出现在 PointPillarScatter 模块中：
def forward(self, voxel_features, coors, batch_size=None):
"""Foraward function to scatter features."""
# TODO: rewrite the function in a batch manner
# no need to deal with different batch cases
if batch_size is not None:
return self.forward_batch(voxel_features, coors, batch_size)
else:
return self.forward_single(voxel_features, coors)
def forward_batch(self, voxel_features, coors, batch_size):
batch_canvas = []
for batch_itt in range(batch_size):
# Create the canvas for this sample
canvas = torch.zeros(
self.in_channels,
self.nx * self.ny,
dtype=voxel_features.dtype,
device=voxel_features.device)
# Only include non-empty pillars
batch_mask = coors[:, 0] == batch_itt
this_coors = coors[batch_mask, :] # 这里产生了 nonzero 算子
... 我们发现问题出现在从输入中拆分 batch 的时候，但是我们实际推理使用时并不需要 batch 维度，即默认 batch=1，所以我们跳过这个操作，直接将输入的 voxel feauters 展开：
def forward(self, voxel_features, coors, batch_size=None):
canvas = torch.zeros(
self.in_channels,
self.nx * self.ny,
dtype=voxel_features.dtype,
device=voxel_features.device)
indices = coors[:, 2] * self.nx + coors[:, 3]
indices = indices.long()
voxels = voxel_features.t()
# Now scatter the blob back to the canvas.
canvas[:, indices] = voxels
# Undo the column stacking to final 4-dim tensor
canvas = canvas.view(1, self.in_channels, self.ny, self.nx)
return canvas 这样我们再尝试一下推理并可视化推理结果，发现已经可以输出正确的结果了，如下图所示：
总结
总的来说 PointPillars 是一个既简单又实用的模型，在保持较高精度的同时又有很高的推理速度，同时部署也很友好，是一个十分常用的模型。目前 MMDeploy 中已经支持了 MMDetection3D 的 PointPillars 和 CenterPoint（pillar 版本）的一键模型转换部署，欢迎大家使用 MMDeploy 来快速地部署 MMDetection3D 的模型。
https://github.com/open-mmlab/mmdetection3dgithub.com/open-mmlab/mmdetection3d
https://github.com/open-mmlab/mmdeploygithub.com/open-mmlab/mmdeploy
码字不易，觉得有用，欢迎点赞、收藏、转发哦，比心~
