[paper review] UniTR : A Unified and Efficient Multi-Modal Transformer for BEV Representation 논문 리뷰

안녕하세요. 이번 포스팅은 UniTR이라는 multi-sensor(camera,lidar) detector에 대해 포스팅하겠습니다.

앞서 아래의 BEVfusion과 DSVT논문을 먼저 읽는 것을 추천합니다.

[paper review] BEVFusion 논문 리뷰

[paper review] DSVT: Dynamic Sparse Voxel Transformer with Rotated Sets 논문리뷰

 

Introduction

논문은 우선 modality-agnostic transformer encoder를 강조합니다. 기존의 모델들은 각각 sensor data마다 encoding layer를 거쳐서 후에 fusion을 한다고 하면 UniTR은 encoding할 때 부터 multi-sensor들을 조합하여 encoding하게 됩니다. 이를 통해 속도 및 성능이 이점이 있다고 후에 ablation study에서 기술하고 있습니다. 그리고 Unify하여 encoding하는 layer의 아키텍쳐가 바로 본 논문에서 소개하는 UniTR이고 pointcloud detector인 DSVT의 transformer의 block과 거의 유사한 흐름을 띄고 있습니다.

 

---

Single-Modal Representation Learning

첫번째 module에서는 각각의 sensor(여기서는 multi-cam, lidar)에서 들어온 raw-data를 tokenization과정을 거칩니다. lidar data는 DSVT에서와 같이 voxelization하여 voxel 내부의 raw-data를 dynamic 하게 encoding하는 dynamic VFE과정을 거치고 image의 경우 ViT에서 사용한 image patch tokenizer방법을 사용합니다. 

class DynamicVoxelVFE(VFETemplate):
    def __init__(self, model_cfg, num_point_features, voxel_size, grid_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

        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(
                PFNLayerV2(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]

        self.scale_xyz = grid_size[0] * grid_size[1] * grid_size[2]
        self.scale_yz = grid_size[1] * grid_size[2]
        self.scale_z = grid_size[2]

        self.grid_size = torch.tensor(grid_size).cuda()
        self.voxel_size = torch.tensor(voxel_size).cuda()
        self.point_cloud_range = torch.tensor(point_cloud_range).cuda()

    def get_output_feature_dim(self):
        return self.num_filters[-1]

    def forward(self, batch_dict, **kwargs):
        points = batch_dict['points'] # (batch_idx, x, y, z, i, e)

        points_coords = torch.floor((points[:, [1,2,3]] - self.point_cloud_range[[0,1,2]]) / self.voxel_size[[0,1,2]]).int()
        mask = ((points_coords >= 0) & (points_coords < self.grid_size[[0,1,2]])).all(dim=1)
        points = points[mask]
        points_coords = points_coords[mask]
        points_xyz = points[:, [1, 2, 3]].contiguous()

        merge_coords = points[:, 0].int() * self.scale_xyz + \
                       points_coords[:, 0] * self.scale_yz + \
                       points_coords[:, 1] * self.scale_z + \
                       points_coords[:, 2]

        unq_coords, unq_inv, unq_cnt = torch.unique(merge_coords, return_inverse=True, return_counts=True, dim=0)

        points_mean = torch_scatter.scatter_mean(points_xyz, unq_inv, dim=0)
        f_cluster = points_xyz - points_mean[unq_inv, :]

        f_center = torch.zeros_like(points_xyz)
        f_center[:, 0] = points_xyz[:, 0] - (points_coords[:, 0].to(points_xyz.dtype) * self.voxel_x + self.x_offset)
        f_center[:, 1] = points_xyz[:, 1] - (points_coords[:, 1].to(points_xyz.dtype) * self.voxel_y + self.y_offset)
        # f_center[:, 2] = points_xyz[:, 2] - self.z_offset
        f_center[:, 2] = points_xyz[:, 2] - (points_coords[:, 2].to(points_xyz.dtype) * self.voxel_z + self.z_offset)

        if self.use_absolute_xyz:
            features = [points[:, 1:], f_cluster, f_center]
        else:
            features = [points[:, 4:], f_cluster, f_center]

        if self.with_distance:
            points_dist = torch.norm(points[:, 1:4], 2, dim=1, keepdim=True)
            features.append(points_dist)
        features = torch.cat(features, dim=-1)

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

        # generate voxel coordinates
        unq_coords = unq_coords.int()
        voxel_coords = torch.stack((unq_coords // self.scale_xyz,
                                    (unq_coords % self.scale_xyz) // self.scale_yz,
                                    (unq_coords % self.scale_yz) // self.scale_z,
                                    unq_coords % self.scale_z), dim=1)
        voxel_coords = voxel_coords[:, [0, 3, 2, 1]]

        batch_dict['pillar_features'] = batch_dict['voxel_features'] = features
        batch_dict['voxel_coords'] = voxel_coords

        return batch_dict

- details

image : 6x256x704x3 (b,h,w,c) ->linear embedding -> m x c 개의 patch로 변경

pointcloud : voxel size (0.3, 0.3, 8.0)m ->Dynamic VFE -> n x c 개의 patch로 변경

multi-sensor patch : (nxm) x c

 

tokenization된 vector들은 위의 그림에서 첫번째인 Intra-modal Block의 input으로 흘러가서 self-attention과정을 거치게 됩니다.. 

구체적으로 수식으로 살펴보면,

첫번째 수식이 의미하는 바는 P(pointcloud)의 경우 총 n개의 token마다 해당 token의 lidar coordinate상의 position(x,y,z)와 feature를 가지고 있고, I(image)의 경우 총 m개의 token마다 image plane상의 x,y position b의 경우 multi-cam의 ID를 의미합니다. x,y는 아마 기준 image(cam)으로 calibration된 위치라고 추측됩니다. 

두번째 수식의 DSP는 DSVT에서의 block을 뜻하며  위의 그림에서 intra-modal Block을 의미합니다. 즉, non-overlap window내에서 위의 주어진 n개만큼 x axis, y axis기준으로 sort된 순서로 set을 만듭니다. pointcloud의 경우 non-empty한 token만을 담습니다. 이렇게 생성된 set을 self-attention과정을 거쳐서 input과 동일한 feature vector가 나오게 됩니다.

마지막 수식은 self-attention과정을 의미하며 MHSA가 multi-head attention을 의미하고 PE는 positional encoding을 의미합니다.

---

Cross-Modal Representation Learning

 

 

다음의 sequence는 inter-modal Block을 거치게 되는데 의미하는 바는 multi-sensor의 경우(위 경우 lidar,camera)센서간의 view-discrepant가 발생하기 때문에 이를 보완하기 위한 block이다. 기존의 모델들은 각각 sensor마다 DL model들을 만들고 이를 통해 나온 feature들을 BEV, 3D or 2D plane에서 fusion하는 방식이지만 본 논문은 하나의 model에서 여러 sensor의 feature를 뽑을 수 있어서 시간적으로 굉장한 이점이 있다고 말하고 있습니다.

먼저 Image space를 기준인 inter-model block을 거치게 되는데 우선 모든 lidar token을 calibration parameter를 이용하여 image plane으로 projection을 합니다. projection해서 매칭되는 image token과는 unify하는데 구체적으로 논문에서 어떻게하는지는 설명되지 않았지만 동일 채널을 가지고 있기에 sum하고 완벽하게 매칭되지 않는 부분은 interpolation할 것으로 추측됩니다.

이렇게 unify된 token들은 위의 block과 동일하게 window 내에서 x axis, y axis방향으로 sort한 후에 set을 만들고 그 set끼리 self-attention을 거쳐서 input과 동일한 shape의 output을 만들게 됩니다.

다음 module은 Inter-modal Block의 3D로 lidar plane에서 encoding과정을 거치는 부분입니다. 앞서와 반대로 이번에는 image token을 lidar plane으로 projection하게 됩니다. 문제는 depth scale이 ambiguity하기 때문에 ill-posed problem이 발생합니다. 이를 해결하기 위해 MVP의 논문의 아이디어를 차용하였습니다. 

 

간략하게 설명하면, 기존의 pointcloud를 2d image plane에 projection해서 image token과 pointcloud 매칭하고 다시 image token을 3D plane으로 projection할때 매칭했던 depth를retrieval합니다. 여기에 매칭되지 않은 image token은 가장 가까운위치의 depth를 retrieve하고 그렇게 3D로 projection된 image token은(retrival 할수없어서 nearest depth로 retrieve한 token)원래 매칭되는 image token과의 픽셀 차이(intensity)를 mlp에 태워 feature에 sum하여 보정하게 됩니다.

 

이후에는 이전과정과 동일한 Block을 거치게 됩니다.

unify된 token들은 window 내에서 x axis, y axis방향으로 sort한 후에 set을 만들고 그 set끼리 self-attention을 거쳐서 input과 동일한 shape의 output을 만들게 됩니다. (x axis로 sort후 MHA, y axis로 sort후 MHA입니다.)

Perception Task

UniTR은 head에 detection head 혹은 segmentation head를 추가하느냐에 따라 3D detection, 3D segmentation task를 수행할 수 있는 versatility한 model이라고 말합니다. 실험에서는 detection의 경우 BEVplane으로 projection하여 (BEVFusion에서 사용한 BEVpooling을 사용) detection을 수행하였습니다.

appendix에 나와있는 ablation study인데, 위의 3가지 종류의 block을 몇개씩 있을때 가장 성능이 좋았는지의 실험결과입니다.

Experiment

 

실험결과는 nuscenes dataset의 결과입니다. 가장 높은 performance결과와 동시에 매우 빠른 inference time을 나타냅니다.  위의 표는 detector model의 벤치마크, 아래표는 segmentation model의 벤치마크입니다. 

LSS는 BEVfusion에서나온 방법으로 image feature를 BEV Plane으로 매핑하는 방법입니다.

특이할만한 실험은 위의 의도적으로 failure case를 만들었을때도 즉, sensor가 malfunction할때도 굉장히 robust한 결과를 inference한다는 점 입니다.

UniTR은 결론적으로 컴팩트하면서 높은 성능을 자랑하며 동시에  standard modulde로 이뤄져있어 비교적 쉽게 deployment하고 빠른 inference time을 갖는 매우 괜찮은 모델이라는 생각이 듭니다.