本文pointpillars基于OpenPcDet框架，转出一个完整的onnx，本文转的onnx是静态的，也可以将voxel个数改为动态输入
注意：推理要使用TensoRT8以上，我用的最新的TensoRT8.4.3.1，不然转出SctterND TensoRT不支持

转一个onnx重点：scatter部分需要改写，将VFE和Scatter两部分合并，代码如下：

class PillarVFE(VFETemplate):
    def __init__(self, model_cfg, num_point_features, voxel_size, point_cloud_range, **kwargs):
        super().__init__(model_cfg=model_cfg)

        self.use_norm = self.model_cfg.USE_NORM
        self.with_distance = self.model_cfg.WITH_DISTANCE
        self.use_absolute_xyz = self.model_cfg.USE_ABSLOTE_XYZ

        num_point_features += 6 if self.use_absolute_xyz else 3
        if self.with_distance:
            num_point_features += 1
        # if self.model_cfg.WITH_NUM_POINTS:
        #     num_point_features += 1
        self.num_point_features = num_point_features
        self.num_filters = self.model_cfg.NUM_FILTERS
        assert len(self.num_filters) > 0
        num_filters = [num_point_features] + list(self.num_filters)

        pfn_layers = []
        for i in range(len(num_filters) - 1):
            in_filters = num_filters[i]
            out_filters = num_filters[i + 1]
            pfn_layers.append(
                PFNLayer(in_filters, out_filters, self.use_norm, last_layer=(i >= len(num_filters) - 2))
            )
        self.pfn_layers = nn.ModuleList(pfn_layers)

        self.voxel_x = voxel_size[0]
        self.voxel_y = voxel_size[1]
        self.voxel_z = voxel_size[2]
        self.x_offset = self.voxel_x / 2 + point_cloud_range[0]
        self.y_offset = self.voxel_y / 2 + point_cloud_range[1]
        self.z_offset = self.voxel_z / 2 + point_cloud_range[2]

    def forward(self, voxel_features, voxel_num_points, coords):
        features_ls = [voxel_features]
        print(voxel_num_points.view(-1, 1, 1).shape) # torch.Size([20000, 1, 1])
        points_mean = voxel_features[:, :, :3].sum(dim=1, keepdim=True) / voxel_num_points.type_as(voxel_features).view(-1, 1, 1)
        # points_mean = voxel_features[:, :, :3].sum(dim=1, keepdim=True) / voxel_num_points.type_as(voxel_features).view(-1, 1, 1)
        f_cluster = voxel_features[:, :, :3] - points_mean
        features_ls.append(f_cluster)

        device = voxel_features.device
        # 使用矩阵切片代替取下表方式
        f_center = voxel_features[..., :3] - (coords[..., 1:] * torch.tensor([self.voxel_z, self.voxel_y, self.voxel_x]).to(device) + torch.tensor([self.z_offset, self.y_offset, self.x_offset]).to(device)).unsqueeze(1).flip(2)
        
        features_ls.append(f_center)
        features = torch.cat(features_ls, dim=-1)

        voxel_count = features.shape[1] # 32
        mask = get_paddings_indicator(voxel_num_points, voxel_count, axis=0)
        mask = torch.unsqueeze(mask, -1).type_as(features)
        features *= mask

        for pfn in self.pfn_layers:
            features = pfn(features)

        # float32[voxel_num,1,64]第一维度为动态的，需要指定压缩维度
        features = torch.squeeze(features, 1)  # float32[20000,64]
        pillars_feature = features.t()  # float32[64,20000]

        spatial_feature = torch.zeros(64, 432 * 496,dtype=features.dtype, device=features.device)

        indices =  coords[:, 2] * 432 + coords[:, 3] #432 * y + x
        indices = indices.long()

        spatial_feature[:, indices] = pillars_feature
        spatial_feature = spatial_feature.view(1,64, 496, 432) # 对应onnx resahap

        return torch.where(torch.isnan(spatial_feature), torch.full_like(spatial_feature, 0), spatial_feature)

在torch和onnx精度对齐过程中发现，经过scatterND reshape节点后，输出值中有部分为nan（应该是输入的点存在nan，没有处理），影响后续计算，决定认为将输出值为nan替换为0,可以在onnx图中看出多了isnan和where节点

以mse作为torch与onnx精度评估：

        cls_preds_mse = np.mean(outputs[0] - cls_preds.cpu().numpy())**2
        box_preds_mse = np.mean(outputs[1] - box_preds.cpu().numpy())**2
        dir_preds_mse = np.mean(outputs[2] - dir_preds.cpu().numpy())**2

        print("cls_preds_mse :", cls_preds_mse)  # 1.0237494843803958e-09
        print("box_preds_mse :", box_preds_mse)  # 6.773186036880116e-13
        print("dir_preds_mse :",dir_preds_mse)  # 2.2815304326538114e-15

看误差，转出onnx应该没是问题的