add code

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/bbox/coders/centerpoint_bbox_coders.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+
+from mmdet.core.bbox import BaseBBoxCoder
+from mmdet.core.bbox.builder import BBOX_CODERS
+
+
+@BBOX_CODERS.register_module(force=True)
+class CenterPointBBoxCoder(BaseBBoxCoder):
+    """Bbox coder for CenterPoint.
+
+    Args:
+        pc_range (list[float]): Range of point cloud.
+        out_size_factor (int): Downsample factor of the model.
+        voxel_size (list[float]): Size of voxel.
+        post_center_range (list[float], optional): Limit of the center.
+            Default: None.
+        max_num (int, optional): Max number to be kept. Default: 100.
+        score_threshold (float, optional): Threshold to filter boxes
+            based on score. Default: None.
+        code_size (int, optional): Code size of bboxes. Default: 9
+    """
+
+    def __init__(self,
+                 pc_range,
+                 out_size_factor,
+                 voxel_size,
+                 post_center_range=None,
+                 max_num=100,
+                 score_threshold=None,
+                 code_size=9):
+
+        self.pc_range = pc_range    # [x_min, y_min, ...]
+        self.out_size_factor = out_size_factor
+        self.voxel_size = voxel_size
+        self.post_center_range = post_center_range  # [-61.2, -61.2, -10.0, 61.2, 61.2, 10.0]
+        self.max_num = max_num
+        self.score_threshold = score_threshold
+        self.code_size = code_size
+
+    def _gather_feat(self, feats, inds, feat_masks=None):
+        """Given feats and indexes, returns the gathered feats.
+
+        Args:
+            feats (torch.Tensor): Features to be transposed and gathered
+                with the shape of [B, 2, W, H].
+            inds (torch.Tensor): Indexes with the shape of [B, N].
+            feat_masks (torch.Tensor, optional): Mask of the feats.
+                Default: None.
+
+        Returns:
+            torch.Tensor: Gathered feats.
+        """
+        dim = feats.size(2)
+        inds = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), dim)
+        feats = feats.gather(1, inds)
+        if feat_masks is not None:
+            feat_masks = feat_masks.unsqueeze(2).expand_as(feats)
+            feats = feats[feat_masks]
+            feats = feats.view(-1, dim)
+        return feats
+
+    def _topk(self, scores, K=80):
+        """Get indexes based on scores.
+
+        Args:
+            scores (torch.Tensor): scores with the shape of (B, N_cls, H, W).
+            K (int, optional): Number to be kept. Defaults to 80.
+
+        Returns:
+            tuple[torch.Tensor]
+                torch.Tensor: Selected scores with the shape of [B, K].
+                torch.Tensor: Selected indexes with the shape of [B, K].
+                torch.Tensor: Selected classes with the shape of [B, K].
+                torch.Tensor: Selected y coord with the shape of [B, K].
+                torch.Tensor: Selected x coord with the shape of [B, K].
+        """
+        batch, cat, height, width = scores.size()
+
+        # 先是针对每个类别的预测都取topK.
+        # (B, N_cls, K), (B, N_cls, K)
+        topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)
+
+        topk_inds = topk_inds % (height * width)    # (B, N_cls, K), topK对应的像素索引(0, H*W-1).
+        topk_ys = (topk_inds.float() /
+                   torch.tensor(width, dtype=torch.float)).int().float()    # (B, N_cls, K), y坐标.
+        topk_xs = (topk_inds % width).int().float()     # (B, N_cls, K), x坐标.
+
+        # 然后对将所有类别得到的topK数据再次进行topK.
+        # (B, K), (B, K)
+        topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
+        topk_clses = (topk_ind / torch.tensor(K, dtype=torch.float)).int()      # (B, K)  对应的类别.
+        # (B, N_cls*K, 1) --gather--> (B, K, 1) --> (B, K)  topK对应的像素坐标索引(0, H*W-1).
+        topk_inds = self._gather_feat(topk_inds.view(batch, -1, 1),
+                                      topk_ind).view(batch, K)
+        # (B, N_cls*K, 1) --gather--> (B, K, 1) --> (B, K)  topK对应的y坐标.
+        topk_ys = self._gather_feat(topk_ys.view(batch, -1, 1),
+                                    topk_ind).view(batch, K)
+        # (B, N_cls*K, 1) --gather--> (B, K, 1) --> (B, K)  topK对应的x坐标.
+        topk_xs = self._gather_feat(topk_xs.view(batch, -1, 1),
+                                    topk_ind).view(batch, K)
+
+        return topk_score, topk_inds, topk_clses, topk_ys, topk_xs
+
+    def _transpose_and_gather_feat(self, feat, ind):
+        """Given feats and indexes, returns the transposed and gathered feats.
+
+        Args:
+            feat (torch.Tensor): Features to be transposed and gathered
+                with the shape of (B, N_c, H, W).
+            ind (torch.Tensor): Indexes with the shape of [B, K].
+
+        Returns:
+            torch.Tensor: Transposed and gathered feats.
+        """
+        # (B, N_c, H, W) --> (B, H, W, N_c) --> (B, H*W, N_c)
+        feat = feat.permute(0, 2, 3, 1).contiguous()
+        feat = feat.view(feat.size(0), -1, feat.size(3))
+        feat = self._gather_feat(feat, ind)     # (B, K, N_c)
+        return feat
+
+    def encode(self):
+        pass
+
+    def decode(self,
+               heat,
+               rot_sine,
+               rot_cosine,
+               hei,
+               dim,
+               vel,
+               reg=None,
+               task_id=-1):
+        """Decode bboxes.
+
+        Args:
+            heat (torch.Tensor): Heatmap with the shape of (B, N_cls, H, W).
+            rot_sine (torch.Tensor): Sine of rotation with the shape of (B, 1, H, W).
+            rot_cosine (torch.Tensor): Cosine of rotation with the shape of (B, 1, H, W).
+            hei (torch.Tensor): Height of the boxes with the shape of (B, 1, H, W).
+            dim (torch.Tensor): Dim of the boxes with the shape of (B, 3, H, W).
+            vel (torch.Tensor): Velocity with the shape of (B, 1, H, W).
+            reg (torch.Tensor, optional): Regression value of the boxes in
+                2D with the shape of (B, 2, H, W). Default: None.
+            task_id (int, optional): Index of task. Default: -1.
+
+        Returns:
+            list[dict]: Decoded boxes.  List[p_dict0, p_dict1, ...]
+                p_dict = {
+                    'bboxes': boxes3d,      # (K', 9)
+                    'scores': scores,       # (K', )
+                    'labels': labels        # (K', )
+                }
+        """
+        batch, cat, _, _ = heat.size()
+
+        # (B, K)
+        scores, inds, clses, ys, xs = self._topk(heat, K=self.max_num)
+
+        if reg is not None:
+            reg = self._transpose_and_gather_feat(reg, inds)    # (B, K, 2)
+            reg = reg.view(batch, self.max_num, 2)
+            xs = xs.view(batch, self.max_num, 1) + reg[:, :, 0:1]    # (B, K, 1) + (B, K, 1) --> (B, K, 1)
+            ys = ys.view(batch, self.max_num, 1) + reg[:, :, 1:2]    # (B, K, 1) + (B, K, 1) --> (B, K, 1)
+        else:
+            xs = xs.view(batch, self.max_num, 1) + 0.5
+            ys = ys.view(batch, self.max_num, 1) + 0.5
+
+        # rotation value and direction label
+        rot_sine = self._transpose_and_gather_feat(rot_sine, inds)   # (B, K, 1)
+        rot_sine = rot_sine.view(batch, self.max_num, 1)
+
+        rot_cosine = self._transpose_and_gather_feat(rot_cosine, inds)  # (B, K, 1)
+        rot_cosine = rot_cosine.view(batch, self.max_num, 1)
+        rot = torch.atan2(rot_sine, rot_cosine)     # (B, K, 1)
+
+        # height in the bev
+        hei = self._transpose_and_gather_feat(hei, inds)
+        hei = hei.view(batch, self.max_num, 1)      # (B, K, 1)
+
+        # dim of the box
+        dim = self._transpose_and_gather_feat(dim, inds)
+        dim = dim.view(batch, self.max_num, 3)      # (B, K, 3)
+
+        # class label
+        clses = clses.view(batch, self.max_num).float()     # (B, K)
+        scores = scores.view(batch, self.max_num)   # (B, K)
+
+        # 计算真实的bev坐标.
+        xs = xs.view(
+            batch, self.max_num,
+            1) * self.out_size_factor * self.voxel_size[0] + self.pc_range[0]
+        ys = ys.view(
+            batch, self.max_num,
+            1) * self.out_size_factor * self.voxel_size[1] + self.pc_range[1]
+
+        if vel is None:  # KITTI FORMAT
+            final_box_preds = torch.cat([xs, ys, hei, dim, rot], dim=2)
+        else:  # exist velocity, nuscene format
+            vel = self._transpose_and_gather_feat(vel, inds)    # (B, K, 2)
+            vel = vel.view(batch, self.max_num, 2)
+            final_box_preds = torch.cat([xs, ys, hei, dim, rot, vel], dim=2)    # (B, K, 9)
+
+        final_scores = scores
+        final_preds = clses
+
+        # use score threshold
+        if self.score_threshold is not None:
+            thresh_mask = final_scores > self.score_threshold   # (B, K)
+
+        if self.post_center_range is not None:
+            self.post_center_range = torch.tensor(
+                self.post_center_range, device=heat.device)
+            mask = (final_box_preds[..., :3] >=
+                    self.post_center_range[:3]).all(2)      # (B, K, 3) --> (B, K)
+            mask &= (final_box_preds[..., :3] <=
+                     self.post_center_range[3:]).all(2)     # (B, K, 3) --> (B, K)
+
+            predictions_dicts = []
+            for i in range(batch):
+                cmask = mask[i, :]      # (K, )
+                if self.score_threshold:
+                    cmask &= thresh_mask[i]     # (K, )
+
+                boxes3d = final_box_preds[i, cmask]     # (K', 9)
+                scores = final_scores[i, cmask]         # (K', )
+                labels = final_preds[i, cmask]          # (K', )
+                predictions_dict = {
+                    'bboxes': boxes3d,      # (K', 9)
+                    'scores': scores,       # (K', )
+                    'labels': labels        # (K', )
+                }
+
+                # List[p_dict0, p_dict1, ...]   len = batch_size
+                predictions_dicts.append(predictions_dict)
+        else:
+            raise NotImplementedError(
+                'Need to reorganize output as a batch, only '
+                'support post_center_range is not None for now!')
+
+        return predictions_dicts
+
+    def center_decode(self,
+               heat,
+               hei,
+               reg=None,
+               task_id=-1):
+        batch, cat, _, _ = heat.size()
+
+        # (B, K)
+        scores, inds, clses, ys, xs = self._topk(heat, K=self.max_num)
+
+        if reg is not None:
+            reg = self._transpose_and_gather_feat(reg, inds)    # (B, K, 2)
+            reg = reg.view(batch, self.max_num, 2)
+            xs = xs.view(batch, self.max_num, 1) + reg[:, :, 0:1]    # (B, K, 1) + (B, K, 1) --> (B, K, 1)
+            ys = ys.view(batch, self.max_num, 1) + reg[:, :, 1:2]    # (B, K, 1) + (B, K, 1) --> (B, K, 1)
+        else:
+            xs = xs.view(batch, self.max_num, 1) + 0.5
+            ys = ys.view(batch, self.max_num, 1) + 0.5
+
+        # height in the bev
+        hei = self._transpose_and_gather_feat(hei, inds)
+        hei = hei.view(batch, self.max_num, 1)      # (B, K, 1)
+
+        # class label
+        clses = clses.view(batch, self.max_num).float()     # (B, K)
+        scores = scores.view(batch, self.max_num)   # (B, K)
+
+        # 计算真实的bev坐标.
+        xs = xs.view(
+            batch, self.max_num,
+            1) * self.out_size_factor * self.voxel_size[0] + self.pc_range[0]
+        ys = ys.view(
+            batch, self.max_num,
+            1) * self.out_size_factor * self.voxel_size[1] + self.pc_range[1]
+
+        final_center_preds = torch.cat([xs, ys, hei], dim=2)
+        final_scores = scores
+        final_preds = clses
+
+        # use score threshold
+        if self.score_threshold is not None:
+            thresh_mask = final_scores > self.score_threshold   # (B, K)
+
+        if self.post_center_range is not None:
+            self.post_center_range = torch.tensor(
+                self.post_center_range, device=heat.device)
+            mask = (final_center_preds[..., :3] >=
+                    self.post_center_range[:3]).all(2)      # (B, K, 3) --> (B, K)
+            mask &= (final_center_preds[..., :3] <=
+                     self.post_center_range[3:]).all(2)     # (B, K, 3) --> (B, K)
+
+            predictions_dicts = []
+            for i in range(batch):
+                cmask = mask[i, :]      # (K, )
+                if self.score_threshold:
+                    cmask &= thresh_mask[i]     # (K, )
+
+                centers = final_center_preds[i, cmask]     # (K', 9)
+                scores = final_scores[i, cmask]         # (K', )
+                labels = final_preds[i, cmask]          # (K', )
+                predictions_dict = {
+                    'centers': centers,      # (K', 9)
+                    'scores': scores,       # (K', )
+                    'labels': labels        # (K', )
+                }
+
+                # List[p_dict0, p_dict1, ...]   len = batch_size
+                predictions_dicts.append(predictions_dict)
+        else:
+            raise NotImplementedError(
+                'Need to reorganize output as a batch, only '
+                'support post_center_range is not None for now!')
+
+        return predictions_dicts

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/evaluation/occ_metrics.py

+import numpy as np
+import os
+from pathlib import Path
+from tqdm import tqdm
+import pickle as pkl
+import argparse
+import time
+import torch
+import sys, platform
+from sklearn.neighbors import KDTree
+from termcolor import colored
+from pathlib import Path
+from copy import deepcopy
+from functools import reduce
+
+np.seterr(divide='ignore', invalid='ignore')
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+def pcolor(string, color, on_color=None, attrs=None):
+    """
+    Produces a colored string for printing
+
+    Parameters
+    ----------
+    string : str
+        String that will be colored
+    color : str
+        Color to use
+    on_color : str
+        Background color to use
+    attrs : list of str
+        Different attributes for the string
+
+    Returns
+    -------
+    string: str
+        Colored string
+    """
+    return colored(string, color, on_color, attrs)
+
+
+def getCellCoordinates(points, voxelSize):
+    return (points / voxelSize).astype(np.int)
+
+
+def getNumUniqueCells(cells):
+    M = cells.max() + 1
+    return np.unique(cells[:, 0] + M * cells[:, 1] + M ** 2 * cells[:, 2]).shape[0]
+
+
+class Metric_mIoU():
+    def __init__(self,
+                 save_dir='.',
+                 num_classes=18,
+                 use_lidar_mask=False,
+                 use_image_mask=False,
+                 ):
+        self.class_names = ['others','barrier', 'bicycle', 'bus', 'car', 'construction_vehicle',
+                            'motorcycle', 'pedestrian', 'traffic_cone', 'trailer', 'truck',
+                            'driveable_surface', 'other_flat', 'sidewalk',
+                            'terrain', 'manmade', 'vegetation','free']
+        self.save_dir = save_dir
+        self.use_lidar_mask = use_lidar_mask
+        self.use_image_mask = use_image_mask
+        self.num_classes = num_classes
+
+        self.point_cloud_range = [-40.0, -40.0, -1.0, 40.0, 40.0, 5.4]
+        self.occupancy_size = [0.4, 0.4, 0.4]
+        self.voxel_size = 0.4
+        self.occ_xdim = int((self.point_cloud_range[3] - self.point_cloud_range[0]) / self.occupancy_size[0])
+        self.occ_ydim = int((self.point_cloud_range[4] - self.point_cloud_range[1]) / self.occupancy_size[1])
+        self.occ_zdim = int((self.point_cloud_range[5] - self.point_cloud_range[2]) / self.occupancy_size[2])
+        self.voxel_num = self.occ_xdim * self.occ_ydim * self.occ_zdim
+        self.hist = np.zeros((self.num_classes, self.num_classes))
+        self.cnt = 0
+
+    def hist_info(self, n_cl, pred, gt):
+        """
+        build confusion matrix
+        # empty classes:0
+        non-empty class: 0-16
+        free voxel class: 17
+
+        Args:
+            n_cl (int): num_classes_occupancy
+            pred (1-d array): pred_occupancy_label, (N_valid, )
+            gt (1-d array): gt_occupancu_label, (N_valid, )
+
+        Returns:
+            tuple:(hist, correctly number_predicted_labels, num_labelled_sample)
+        """
+        assert pred.shape == gt.shape
+        k = (gt >= 0) & (gt < n_cl)  # exclude 255
+        labeled = np.sum(k)     # N_total
+        correct = np.sum((pred[k] == gt[k]))    # N_correct
+
+        return (
+            np.bincount(
+                n_cl * gt[k].astype(int) + pred[k].astype(int), minlength=n_cl ** 2
+            ).reshape(n_cl, n_cl),    # (N_cls, N_cls),
+            correct,    # N_correct
+            labeled,    # N_total
+        )
+
+    def per_class_iu(self, hist):
+
+        return np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
+
+    def compute_mIoU(self, pred, label, n_classes):
+        """
+        Args:
+            pred: (N_valid, )
+            label: (N_valid, )
+            n_classes: int=18
+
+        Returns:
+
+        """
+        hist = np.zeros((n_classes, n_classes))     # (N_cls, N_cls)
+        new_hist, correct, labeled = self.hist_info(n_classes, pred.flatten(), label.flatten())
+        hist += new_hist    # (N_cls, N_cls)
+        mIoUs = self.per_class_iu(hist)
+        # for ind_class in range(n_classes):
+        #     print(str(round(mIoUs[ind_class] * 100, 2)))
+        # print('===> mIoU: ' + str(round(np.nanmean(mIoUs) * 100, 2)))
+        return round(np.nanmean(mIoUs) * 100, 2), hist
+
+    def add_batch(self, semantics_pred, semantics_gt, mask_lidar, mask_camera):
+        """
+        Args:
+            semantics_pred: (Dx, Dy, Dz, n_cls)
+            semantics_gt: (Dx, Dy, Dz)
+            mask_lidar: (Dx, Dy, Dz)
+            mask_camera: (Dx, Dy, Dz)
+
+        Returns:
+
+        """
+        self.cnt += 1
+        if self.use_image_mask:
+            masked_semantics_gt = semantics_gt[mask_camera]     # (N_valid, )
+            masked_semantics_pred = semantics_pred[mask_camera]     # (N_valid, )
+        elif self.use_lidar_mask:
+            masked_semantics_gt = semantics_gt[mask_lidar]
+            masked_semantics_pred = semantics_pred[mask_lidar]
+        else:
+            masked_semantics_gt = semantics_gt
+            masked_semantics_pred = semantics_pred
+
+            # # pred = np.random.randint(low=0, high=17, size=masked_semantics.shape)
+        _, _hist = self.compute_mIoU(masked_semantics_pred, masked_semantics_gt, self.num_classes)
+        self.hist += _hist    # (N_cls, N_cls)  列对应每个gt类别，行对应每个预测类别, 这样只有对角线位置上的预测是准确的.
+
+    def count_miou(self):
+        mIoU = self.per_class_iu(self.hist)
+        # assert cnt == num_samples, 'some samples are not included in the miou calculation'
+        print(f'===> per class IoU of {self.cnt} samples:')
+        for ind_class in range(self.num_classes-1):
+            print(f'===> {self.class_names[ind_class]} - IoU = ' + str(round(mIoU[ind_class] * 100, 2)))
+
+        print(f'===> mIoU of {self.cnt} samples: ' + str(round(np.nanmean(mIoU[:self.num_classes-1]) * 100, 2)))
+        # print(f'===> sample-wise averaged mIoU of {cnt} samples: ' + str(round(np.nanmean(mIoU_avg), 2)))
+
+        eval_res = dict()
+        # eval_res['class_name'] = self.class_names
+        eval_res['mIoU'] = mIoU
+        # eval_res['cnt'] = self.cnt
+        return eval_res
+
+
+class Metric_FScore():
+    def __init__(self,
+
+                 leaf_size=10,
+                 threshold_acc=0.6,
+                 threshold_complete=0.6,
+                 voxel_size=[0.4, 0.4, 0.4],
+                 range=[-40, -40, -1, 40, 40, 5.4],
+                 void=[17, 255],
+                 use_lidar_mask=False,
+                 use_image_mask=False, ) -> None:
+
+        self.leaf_size = leaf_size
+        self.threshold_acc = threshold_acc
+        self.threshold_complete = threshold_complete
+        self.voxel_size = voxel_size
+        self.range = range
+        self.void = void
+        self.use_lidar_mask = use_lidar_mask
+        self.use_image_mask = use_image_mask
+        self.cnt=0
+        self.tot_acc = 0.
+        self.tot_cmpl = 0.
+        self.tot_f1_mean = 0.
+        self.eps = 1e-8
+
+
+
+    def voxel2points(self, voxel):
+        # occIdx = torch.where(torch.logical_and(voxel != FREE, voxel != NOT_OBSERVED))
+        # if isinstance(voxel, np.ndarray): voxel = torch.from_numpy(voxel)
+        mask = np.logical_not(reduce(np.logical_or, [voxel == self.void[i] for i in range(len(self.void))]))
+        occIdx = np.where(mask)
+
+        points = np.concatenate((occIdx[0][:, None] * self.voxel_size[0] + self.voxel_size[0] / 2 + self.range[0], \
+                                 occIdx[1][:, None] * self.voxel_size[1] + self.voxel_size[1] / 2 + self.range[1], \
+                                 occIdx[2][:, None] * self.voxel_size[2] + self.voxel_size[2] / 2 + self.range[2]),
+                                axis=1)
+        return points
+
+    def add_batch(self, semantics_pred, semantics_gt, mask_lidar, mask_camera ):
+
+        # for scene_token in tqdm(preds_dict.keys()):
+        self.cnt += 1
+
+        if self.use_image_mask:
+            semantics_gt[mask_camera == False] = 255
+            semantics_pred[mask_camera == False] = 255
+        elif self.use_lidar_mask:
+            semantics_gt[mask_lidar == False] = 255
+            semantics_pred[mask_lidar == False] = 255
+        else:
+            pass
+
+        ground_truth = self.voxel2points(semantics_gt)
+        prediction = self.voxel2points(semantics_pred)
+        if prediction.shape[0] == 0:
+            accuracy=0
+            completeness=0
+            fmean=0
+
+        else:
+            prediction_tree = KDTree(prediction, leaf_size=self.leaf_size)
+            ground_truth_tree = KDTree(ground_truth, leaf_size=self.leaf_size)
+            complete_distance, _ = prediction_tree.query(ground_truth)
+            complete_distance = complete_distance.flatten()
+
+            accuracy_distance, _ = ground_truth_tree.query(prediction)
+            accuracy_distance = accuracy_distance.flatten()
+
+            # evaluate completeness
+            complete_mask = complete_distance < self.threshold_complete
+            completeness = complete_mask.mean()
+
+            # evalute accuracy
+            accuracy_mask = accuracy_distance < self.threshold_acc
+            accuracy = accuracy_mask.mean()
+
+            fmean = 2.0 / (1 / (accuracy+self.eps) + 1 / (completeness+self.eps))
+
+        self.tot_acc += accuracy
+        self.tot_cmpl += completeness
+        self.tot_f1_mean += fmean
+
+    def count_fscore(self,):
+        base_color, attrs = 'red', ['bold', 'dark']
+        print(pcolor('\n######## F score: {} #######'.format(self.tot_f1_mean / self.cnt), base_color, attrs=attrs))
+
+

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/evaluation/ray_metrics.py

+# Acknowledgments: https://github.com/tarashakhurana/4d-occ-forecasting
+# Modified by Haisong Liu
+import math
+import copy
+import numpy as np
+import torch
+from torch.utils.cpp_extension import load
+from tqdm import tqdm
+from prettytable import PrettyTable
+from .ray_pq import Metric_RayPQ
+
+
+dvr = load("dvr", sources=["lib/dvr/dvr.cpp", "lib/dvr/dvr.cu"], verbose=True, extra_cuda_cflags=['-allow-unsupported-compiler'])
+_pc_range = [-40, -40, -1.0, 40, 40, 5.4]
+_voxel_size = 0.4
+
+occ_class_names = [
+    'others', 'barrier', 'bicycle', 'bus', 'car', 'construction_vehicle',
+    'motorcycle', 'pedestrian', 'traffic_cone', 'trailer', 'truck',
+    'driveable_surface', 'other_flat', 'sidewalk',
+    'terrain', 'manmade', 'vegetation', 'free'
+]
+
+# https://github.com/tarashakhurana/4d-occ-forecasting/blob/ff986082cd6ea10e67ab7839bf0e654736b3f4e2/test_fgbg.py#L29C1-L46C16
+def get_rendered_pcds(origin, points, tindex, pred_dist):
+    pcds = []
+    
+    for t in range(len(origin)):
+        mask = (tindex == t)
+        # skip the ones with no data
+        if not mask.any():
+            continue
+        _pts = points[mask, :3]
+        # use ground truth lidar points for the raycasting direction
+        v = _pts - origin[t][None, :]
+        d = v / np.sqrt((v ** 2).sum(axis=1, keepdims=True))
+        pred_pts = origin[t][None, :] + d * pred_dist[mask][:, None]
+        pcds.append(torch.from_numpy(pred_pts))
+        
+    return pcds
+
+
+def meshgrid3d(occ_size, pc_range):
+    W, H, D = occ_size
+    
+    xs = torch.linspace(0.5, W - 0.5, W).view(W, 1, 1).expand(W, H, D) / W
+    ys = torch.linspace(0.5, H - 0.5, H).view(1, H, 1).expand(W, H, D) / H
+    zs = torch.linspace(0.5, D - 0.5, D).view(1, 1, D).expand(W, H, D) / D
+    xs = xs * (pc_range[3] - pc_range[0]) + pc_range[0]
+    ys = ys * (pc_range[4] - pc_range[1]) + pc_range[1]
+    zs = zs * (pc_range[5] - pc_range[2]) + pc_range[2]
+    xyz = torch.stack((xs, ys, zs), -1)
+
+    return xyz
+
+
+def generate_lidar_rays():
+    # prepare lidar ray angles
+    pitch_angles = []
+    for k in range(10):
+        angle = math.pi / 2 - math.atan(k + 1)
+        pitch_angles.append(-angle)
+    
+    # nuscenes lidar fov: [0.2107773983152201, -0.5439104895672159] (rad)
+    while pitch_angles[-1] < 0.21:
+        delta = pitch_angles[-1] - pitch_angles[-2]
+        pitch_angles.append(pitch_angles[-1] + delta)
+
+    lidar_rays = []
+    for pitch_angle in pitch_angles:
+        for azimuth_angle in np.arange(0, 360, 1):
+            azimuth_angle = np.deg2rad(azimuth_angle)
+
+            x = np.cos(pitch_angle) * np.cos(azimuth_angle)
+            y = np.cos(pitch_angle) * np.sin(azimuth_angle)
+            z = np.sin(pitch_angle)
+
+            lidar_rays.append((x, y, z))
+
+    return np.array(lidar_rays, dtype=np.float32)
+
+
+def process_one_sample(sem_pred, lidar_rays, output_origin, instance_pred=None):
+    # lidar origin in ego coordinate
+    # lidar_origin = torch.tensor([[[0.9858, 0.0000, 1.8402]]])
+    T = output_origin.shape[1]
+    pred_pcds_t = []
+
+    free_id = len(occ_class_names) - 1 
+    occ_pred = copy.deepcopy(sem_pred)
+    occ_pred[sem_pred < free_id] = 1
+    occ_pred[sem_pred == free_id] = 0
+    occ_pred = occ_pred.permute(2, 1, 0)
+    occ_pred = occ_pred[None, None, :].contiguous().float()
+
+    offset = torch.Tensor(_pc_range[:3])[None, None, :]
+    scaler = torch.Tensor([_voxel_size] * 3)[None, None, :]
+
+    lidar_tindex = torch.zeros([1, lidar_rays.shape[0]])
+    
+    for t in range(T): 
+        lidar_origin = output_origin[:, t:t+1, :]  # [1, 1, 3]
+        lidar_endpts = lidar_rays[None] + lidar_origin  # [1, 15840, 3]
+
+        output_origin_render = ((lidar_origin - offset) / scaler).float()  # [1, 1, 3]
+        output_points_render = ((lidar_endpts - offset) / scaler).float()  # [1, N, 3]
+        output_tindex_render = lidar_tindex  # [1, N], all zeros
+
+        with torch.no_grad():
+            pred_dist, _, coord_index = dvr.render_forward(
+                occ_pred.cuda(),
+                output_origin_render.cuda(),
+                output_points_render.cuda(),
+                output_tindex_render.cuda(),
+                [1, 16, 200, 200],
+                "test"
+            )
+            pred_dist *= _voxel_size
+
+        pred_pcds = get_rendered_pcds(
+            lidar_origin[0].cpu().numpy(),
+            lidar_endpts[0].cpu().numpy(),
+            lidar_tindex[0].cpu().numpy(),
+            pred_dist[0].cpu().numpy()
+        )
+        coord_index = coord_index[0, :, :].long().cpu()  # [N, 3]
+
+        pred_label = sem_pred[coord_index[:, 0], coord_index[:, 1], coord_index[:, 2]][:, None]  # [N, 1]
+        pred_dist = pred_dist[0, :, None].cpu()
+        
+        if instance_pred is not None:
+            pred_instance = instance_pred[coord_index[:, 0], coord_index[:, 1], coord_index[:, 2]][:, None]  # [N, 1]
+            pred_pcds = torch.cat([pred_label.float(), pred_instance.float(), pred_dist], dim=-1)
+        else:
+            pred_pcds = torch.cat([pred_label.float(), pred_dist], dim=-1)
+
+        pred_pcds_t.append(pred_pcds)
+
+    pred_pcds_t = torch.cat(pred_pcds_t, dim=0)
+   
+    return pred_pcds_t.numpy()
+
+
+def calc_metrics(pcd_pred_list, pcd_gt_list):
+    thresholds = [1, 2, 4]
+
+    gt_cnt = np.zeros([len(occ_class_names)])
+    pred_cnt = np.zeros([len(occ_class_names)])
+    tp_cnt = np.zeros([len(thresholds), len(occ_class_names)])
+
+    for pcd_pred, pcd_gt in zip(pcd_pred_list, pcd_gt_list):
+        for j, threshold in enumerate(thresholds):
+            # L1
+            depth_pred = pcd_pred[:, 1]
+            depth_gt = pcd_gt[:, 1]
+            l1_error = np.abs(depth_pred - depth_gt)
+            tp_dist_mask = (l1_error < threshold)
+            
+            for i, cls in enumerate(occ_class_names):
+                cls_id = occ_class_names.index(cls)
+                cls_mask_pred = (pcd_pred[:, 0] == cls_id)
+                cls_mask_gt = (pcd_gt[:, 0] == cls_id)
+
+                gt_cnt_i = cls_mask_gt.sum()
+                pred_cnt_i = cls_mask_pred.sum()
+                if j == 0:
+                    gt_cnt[i] += gt_cnt_i
+                    pred_cnt[i] += pred_cnt_i
+
+                tp_cls = cls_mask_gt & cls_mask_pred  # [N]
+                tp_mask = np.logical_and(tp_cls, tp_dist_mask)
+                tp_cnt[j][i] += tp_mask.sum()
+    
+    iou_list = []
+    for j, threshold in enumerate(thresholds):
+        iou_list.append((tp_cnt[j] / (gt_cnt + pred_cnt - tp_cnt[j]))[:-1])
+
+    return iou_list
+
+
+def main_raypq(sem_pred_list, sem_gt_list, inst_pred_list, inst_gt_list, lidar_origin_list):
+    torch.cuda.empty_cache()
+
+    eval_metrics_pq = Metric_RayPQ(
+        num_classes=len(occ_class_names),
+        thresholds=[1, 2, 4]
+    )
+
+    # generate lidar rays
+    lidar_rays = generate_lidar_rays()
+    lidar_rays = torch.from_numpy(lidar_rays)
+
+    for sem_pred, sem_gt, inst_pred, inst_gt, lidar_origins in \
+        tqdm(zip(sem_pred_list, sem_gt_list, inst_pred_list, inst_gt_list, lidar_origin_list), ncols=50):
+        sem_pred = torch.from_numpy(np.reshape(sem_pred, [200, 200, 16]))
+        sem_gt = torch.from_numpy(np.reshape(sem_gt, [200, 200, 16]))
+
+        inst_pred = torch.from_numpy(np.reshape(inst_pred, [200, 200, 16]))
+        inst_gt = torch.from_numpy(np.reshape(inst_gt, [200, 200, 16]))
+
+        pcd_pred = process_one_sample(sem_pred, lidar_rays, lidar_origins, instance_pred=inst_pred)
+        pcd_gt = process_one_sample(sem_gt, lidar_rays, lidar_origins, instance_pred=inst_gt)
+
+        # evalute on non-free rays
+        valid_mask = (pcd_gt[:, 0].astype(np.int32) != len(occ_class_names) - 1)
+        pcd_pred = pcd_pred[valid_mask]
+        pcd_gt = pcd_gt[valid_mask]
+
+        assert pcd_pred.shape == pcd_gt.shape
+        
+        sem_gt = pcd_gt[:, 0].astype(np.int32)
+        sem_pred = pcd_pred[:, 0].astype(np.int32)
+
+        instances_gt = pcd_gt[:, 1].astype(np.int32)
+        instances_pred = pcd_pred[:, 1].astype(np.int32)
+
+        # L1
+        depth_gt = pcd_gt[:, 2]
+        depth_pred = pcd_pred[:, 2]
+        l1_error = np.abs(depth_pred - depth_gt)
+
+        eval_metrics_pq.add_batch(sem_pred, sem_gt, instances_pred, instances_gt, l1_error)
+
+    torch.cuda.empty_cache()
+
+    return eval_metrics_pq.count_pq()
+
+
+def main(sem_pred_list, sem_gt_list, lidar_origin_list):
+    torch.cuda.empty_cache()
+
+    # generate lidar rays
+    lidar_rays = generate_lidar_rays()
+    lidar_rays = torch.from_numpy(lidar_rays)
+
+    pcd_pred_list, pcd_gt_list = [], []
+    for sem_pred, sem_gt, lidar_origins in tqdm(zip(sem_pred_list, sem_gt_list, lidar_origin_list), ncols=50):
+        sem_pred = torch.from_numpy(np.reshape(sem_pred, [200, 200, 16]))
+        sem_gt = torch.from_numpy(np.reshape(sem_gt, [200, 200, 16]))
+
+        pcd_pred = process_one_sample(sem_pred, lidar_rays, lidar_origins)
+        pcd_gt = process_one_sample(sem_gt, lidar_rays, lidar_origins)
+
+        # evalute on non-free rays
+        valid_mask = (pcd_gt[:, 0].astype(np.int32) != len(occ_class_names) - 1)
+        pcd_pred = pcd_pred[valid_mask]
+        pcd_gt = pcd_gt[valid_mask]
+
+        assert pcd_pred.shape == pcd_gt.shape
+        pcd_pred_list.append(pcd_pred)
+        pcd_gt_list.append(pcd_gt)
+
+    iou_list = calc_metrics(pcd_pred_list, pcd_gt_list)
+    rayiou = np.nanmean(iou_list)
+    rayiou_0 = np.nanmean(iou_list[0])
+    rayiou_1 = np.nanmean(iou_list[1])
+    rayiou_2 = np.nanmean(iou_list[2])
+    
+    table = PrettyTable([
+        'Class Names',
+        'RayIoU@1', 'RayIoU@2', 'RayIoU@4'
+    ])
+    table.float_format = '.3'
+
+    for i in range(len(occ_class_names) - 1):
+        table.add_row([
+            occ_class_names[i],
+            iou_list[0][i], iou_list[1][i], iou_list[2][i]
+        ], divider=(i == len(occ_class_names) - 2))
+    
+    table.add_row(['MEAN', rayiou_0, rayiou_1, rayiou_2])
+
+    print(table)
+
+    torch.cuda.empty_cache()
+
+    return {
+        'RayIoU': rayiou,
+        'RayIoU@1': rayiou_0,
+        'RayIoU@2': rayiou_1,
+        'RayIoU@4': rayiou_2,
+    }

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/evaluation/ray_pq.py

+import numpy as np
+from prettytable import PrettyTable
+
+
+class Metric_RayPQ:
+    def __init__(self, 
+                 num_classes=18,
+                 thresholds=[1, 2, 4]):
+        """
+        Args:
+            ignore_index (llist): Class ids that not be considered in pq counting.
+        """
+        if num_classes == 18:
+            self.class_names = [
+                'others','barrier', 'bicycle', 'bus', 'car', 'construction_vehicle',
+                'motorcycle', 'pedestrian', 'traffic_cone', 'trailer', 'truck',
+                'driveable_surface', 'other_flat', 'sidewalk',
+                'terrain', 'manmade', 'vegetation','free'
+            ]
+        else:
+            raise ValueError
+        
+        self.num_classes = num_classes
+        self.id_offset = 2 ** 16
+        self.eps = 1e-5
+        self.thresholds = thresholds
+        
+        self.min_num_points = 10
+        self.include = np.array(
+            [n for n in range(self.num_classes - 1)],
+            dtype=int)
+        self.cnt = 0
+        
+        # panoptic stuff
+        self.pan_tp = np.zeros([len(self.thresholds), num_classes], dtype=int)
+        self.pan_iou = np.zeros([len(self.thresholds), num_classes], dtype=np.double)
+        self.pan_fp = np.zeros([len(self.thresholds), num_classes], dtype=int)
+        self.pan_fn = np.zeros([len(self.thresholds), num_classes], dtype=int)
+        
+    def add_batch(self,semantics_pred,semantics_gt,instances_pred,instances_gt, l1_error):
+        self.cnt += 1
+        self.add_panoptic_sample(semantics_pred, semantics_gt, instances_pred, instances_gt, l1_error) 
+    
+    def add_panoptic_sample(self, semantics_pred, semantics_gt, instances_pred, instances_gt, l1_error):
+        """Add one sample of panoptic predictions and ground truths for
+        evaluation.
+
+        Args:
+            semantics_pred (np.ndarray): Semantic predictions.
+            semantics_gt (np.ndarray): Semantic ground truths.
+            instances_pred (np.ndarray): Instance predictions.
+            instances_gt (np.ndarray): Instance ground truths.
+        """
+        # get instance_class_id from instance_gt
+        instance_class_ids = [self.num_classes - 1]
+        for i in range(1, instances_gt.max() + 1):
+            class_id = np.unique(semantics_gt[instances_gt == i])
+            # assert class_id.shape[0] == 1, "each instance must belong to only one class"
+            if class_id.shape[0] == 1:
+                instance_class_ids.append(class_id[0])
+            else:
+                instance_class_ids.append(self.num_classes - 1)
+        instance_class_ids = np.array(instance_class_ids)
+
+        instance_count = 1
+        final_instance_class_ids = []
+        final_instances = np.zeros_like(instances_gt)  # empty space has instance id "0"
+
+        for class_id in range(self.num_classes - 1):
+            if np.sum(semantics_gt == class_id) == 0:
+                continue
+
+            if self.class_names[class_id] in ['car', 'truck', 'construction_vehicle', 'bus', 'trailer', 'motorcycle', 'bicycle', 'pedestrian']:
+                # treat as instances
+                for instance_id in range(len(instance_class_ids)):
+                    if instance_class_ids[instance_id] != class_id:
+                        continue
+                    final_instances[instances_gt == instance_id] = instance_count
+                    instance_count += 1
+                    final_instance_class_ids.append(class_id)
+            else:
+                # treat as semantics
+                final_instances[semantics_gt == class_id] = instance_count
+                instance_count += 1
+                final_instance_class_ids.append(class_id)
+                
+        instances_gt = final_instances
+        
+        # avoid zero (ignored label)
+        instances_pred = instances_pred + 1
+        instances_gt = instances_gt + 1
+        
+        for j, threshold in enumerate(self.thresholds):
+            tp_dist_mask = l1_error < threshold
+            # for each class (except the ignored ones)
+            for cl in self.include:
+                # get a class mask
+                pred_inst_in_cl_mask = semantics_pred == cl
+                gt_inst_in_cl_mask = semantics_gt == cl
+
+                # get instance points in class (makes outside stuff 0)
+                pred_inst_in_cl = instances_pred * pred_inst_in_cl_mask.astype(int)
+                gt_inst_in_cl = instances_gt * gt_inst_in_cl_mask.astype(int)
+
+                # generate the areas for each unique instance prediction
+                unique_pred, counts_pred = np.unique(
+                    pred_inst_in_cl[pred_inst_in_cl > 0], return_counts=True)
+                id2idx_pred = {id: idx for idx, id in enumerate(unique_pred)}
+                matched_pred = np.array([False] * unique_pred.shape[0])
+
+                # generate the areas for each unique instance gt_np
+                unique_gt, counts_gt = np.unique(
+                    gt_inst_in_cl[gt_inst_in_cl > 0], return_counts=True)
+                id2idx_gt = {id: idx for idx, id in enumerate(unique_gt)}
+                matched_gt = np.array([False] * unique_gt.shape[0])
+
+                # generate intersection using offset
+                valid_combos = np.logical_and(pred_inst_in_cl > 0,
+                                            gt_inst_in_cl > 0)
+                # add dist_mask
+                valid_combos = np.logical_and(valid_combos, tp_dist_mask)
+
+                id_offset_combo = pred_inst_in_cl[
+                    valid_combos] + self.id_offset * gt_inst_in_cl[valid_combos]
+                unique_combo, counts_combo = np.unique(
+                    id_offset_combo, return_counts=True)
+
+                # generate an intersection map
+                # count the intersections with over 0.5 IoU as TP
+                gt_labels = unique_combo // self.id_offset
+                pred_labels = unique_combo % self.id_offset
+                gt_areas = np.array([counts_gt[id2idx_gt[id]] for id in gt_labels])
+                pred_areas = np.array(
+                    [counts_pred[id2idx_pred[id]] for id in pred_labels])
+                intersections = counts_combo
+                unions = gt_areas + pred_areas - intersections
+                ious = intersections.astype(float) / unions.astype(float)
+
+                tp_indexes = ious > 0.5
+                self.pan_tp[j][cl] += np.sum(tp_indexes)
+                self.pan_iou[j][cl] += np.sum(ious[tp_indexes])
+
+                matched_gt[[id2idx_gt[id] for id in gt_labels[tp_indexes]]] = True
+                matched_pred[[id2idx_pred[id]
+                            for id in pred_labels[tp_indexes]]] = True
+
+                # count the FN
+                if len(counts_gt) > 0:
+                    self.pan_fn[j][cl] += np.sum(
+                        np.logical_and(counts_gt >= self.min_num_points,
+                                    ~matched_gt))
+
+                # count the FP
+                if len(matched_pred) > 0:
+                    self.pan_fp[j][cl] += np.sum(
+                        np.logical_and(counts_pred >= self.min_num_points,
+                                    ~matched_pred))
+    
+    def count_pq(self):
+        sq_all = self.pan_iou.astype(np.double) / np.maximum(
+            self.pan_tp.astype(np.double), self.eps)
+        rq_all = self.pan_tp.astype(np.double) / np.maximum(
+            self.pan_tp.astype(np.double) + 0.5 * self.pan_fp.astype(np.double)
+            + 0.5 * self.pan_fn.astype(np.double), self.eps)
+        pq_all = sq_all * rq_all
+        
+        # mask classes not occurring in dataset
+        mask = (self.pan_tp + self.pan_fp + self.pan_fn) > 0
+        pq_all[~mask] = float('nan')
+
+        table = PrettyTable([
+            'Class Names',
+            'RayPQ@%d' % self.thresholds[0],
+            'RayPQ@%d' % self.thresholds[1],
+            'RayPQ@%d' % self.thresholds[2]
+        ])
+        table.float_format = '.3'
+
+        for i in range(len(self.class_names) - 1):
+            table.add_row([
+                self.class_names[i],
+                pq_all[0][i], pq_all[1][i], pq_all[2][i],
+            ], divider=(i == len(self.class_names) - 2))
+        
+        table.add_row([
+            'MEAN',
+            np.nanmean(pq_all[0]), np.nanmean(pq_all[1]), np.nanmean(pq_all[2])
+        ])
+
+        print(table)
+
+        return {
+            'RayPQ': np.nanmean(pq_all),
+            'RayPQ@1': np.nanmean(pq_all[0]),
+            'RayPQ@2': np.nanmean(pq_all[1]),
+            'RayPQ@4': np.nanmean(pq_all[2]),
+        }

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/hook/__init__.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from .ema import MEGVIIEMAHook
+from .utils import is_parallel
+from .sequentialcontrol import SequentialControlHook
+from .syncbncontrol import SyncbnControlHook
+
+__all__ = ['MEGVIIEMAHook', 'SequentialControlHook', 'is_parallel',
+           'SyncbnControlHook']

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/hook/ema.py

+# Copyright (c) OpenMMLab. All rights reserved.
+# modified from megvii-bevdepth.
+import math
+import os
+from copy import deepcopy
+
+import torch
+from mmcv.runner import load_state_dict
+from mmcv.runner.dist_utils import master_only
+from mmcv.runner.hooks import HOOKS, Hook
+
+from .utils import is_parallel
+
+__all__ = ['ModelEMA']
+
+
+class ModelEMA:
+    """Model Exponential Moving Average from https://github.com/rwightman/
+    pytorch-image-models Keep a moving average of everything in the model
+    state_dict (parameters and buffers).
+
+    This is intended to allow functionality like
+    https://www.tensorflow.org/api_docs/python/tf/train/
+    ExponentialMovingAverage
+    A smoothed version of the weights is necessary for some training
+    schemes to perform well.
+    This class is sensitive where it is initialized in the sequence
+    of model init, GPU assignment and distributed training wrappers.
+    """
+
+    def __init__(self, model, decay=0.9999, updates=0):
+        """
+        Args:
+            model (nn.Module): model to apply EMA.
+            decay (float): ema decay reate.
+            updates (int): counter of EMA updates.
+        """
+        # Create EMA(FP32)
+        self.ema_model = deepcopy(model).eval()
+        self.ema = self.ema_model.module.module if is_parallel(
+            self.ema_model.module) else self.ema_model.module
+        self.updates = updates
+        # decay exponential ramp (to help early epochs)
+        self.decay = lambda x: decay * (1 - math.exp(-x / 2000))
+        for p in self.ema.parameters():
+            p.requires_grad_(False)
+
+    def update(self, trainer, model):
+        # Update EMA parameters
+        with torch.no_grad():
+            self.updates += 1
+            d = self.decay(self.updates)
+
+            msd = model.module.state_dict() if is_parallel(
+                model) else model.state_dict()  # model state_dict
+            for k, v in self.ema.state_dict().items():
+                if v.dtype.is_floating_point:
+                    v *= d
+                    v += (1.0 - d) * msd[k].detach()
+
+
+@HOOKS.register_module()
+class MEGVIIEMAHook(Hook):
+    """EMAHook used in BEVDepth.
+
+    Modified from https://github.com/Megvii-Base
+    Detection/BEVDepth/blob/main/callbacks/ema.py.
+    """
+
+    def __init__(self, init_updates=0, decay=0.9990, resume=None):
+        super().__init__()
+        self.init_updates = init_updates
+        self.resume = resume
+        self.decay = decay
+
+    def before_run(self, runner):
+        from torch.nn.modules.batchnorm import SyncBatchNorm
+
+        bn_model_list = list()
+        bn_model_dist_group_list = list()
+        for model_ref in runner.model.modules():
+            if isinstance(model_ref, SyncBatchNorm):
+                bn_model_list.append(model_ref)
+                bn_model_dist_group_list.append(model_ref.process_group)
+                model_ref.process_group = None
+        runner.ema_model = ModelEMA(runner.model, self.decay)
+
+        for bn_model, dist_group in zip(bn_model_list,
+                                        bn_model_dist_group_list):
+            bn_model.process_group = dist_group
+        runner.ema_model.updates = self.init_updates
+
+        if self.resume is not None:
+            runner.logger.info(f'resume ema checkpoint from {self.resume}')
+            cpt = torch.load(self.resume, map_location='cpu')
+            load_state_dict(runner.ema_model.ema, cpt['state_dict'])
+            runner.ema_model.updates = cpt['updates']
+
+    def after_train_iter(self, runner):
+        runner.ema_model.update(runner, runner.model.module)
+
+    def after_train_epoch(self, runner):
+        # if self.is_last_epoch(runner):   # 只保存最后一个epoch的ema权重.
+        self.save_checkpoint(runner)
+
+    @master_only
+    def save_checkpoint(self, runner):
+        state_dict = runner.ema_model.ema.state_dict()
+        ema_checkpoint = {
+            'epoch': runner.epoch,
+            'state_dict': state_dict,
+            'updates': runner.ema_model.updates
+        }
+        save_path = f'epoch_{runner.epoch+1}_ema.pth'
+        save_path = os.path.join(runner.work_dir, save_path)
+        torch.save(ema_checkpoint, save_path)
+        runner.logger.info(f'Saving ema checkpoint at {save_path}')

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/hook/sequentialcontrol.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from mmcv.runner.hooks import HOOKS, Hook
+from .utils import is_parallel
+
+__all__ = ['SequentialControlHook']
+
+
+@HOOKS.register_module()
+class SequentialControlHook(Hook):
+    """ """
+
+    def __init__(self, temporal_start_epoch=1):
+        super().__init__()
+        self.temporal_start_epoch=temporal_start_epoch
+
+    def set_temporal_flag(self, runner, flag):
+        if is_parallel(runner.model.module):
+            runner.model.module.module.with_prev=flag
+        else:
+            runner.model.module.with_prev = flag
+
+    def before_run(self, runner):
+        self.set_temporal_flag(runner, False)
+
+    def before_train_epoch(self, runner):
+        if runner.epoch > self.temporal_start_epoch:
+            self.set_temporal_flag(runner, True)
\ No newline at end of file

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/hook/syncbncontrol.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from mmcv.runner.hooks import HOOKS, Hook
+from .utils import is_parallel
+from torch.nn import SyncBatchNorm
+
+__all__ = ['SyncbnControlHook']
+
+
+@HOOKS.register_module()
+class SyncbnControlHook(Hook):
+    """ """
+
+    def __init__(self, syncbn_start_epoch=1):
+        super().__init__()
+        self.is_syncbn=False
+        self.syncbn_start_epoch = syncbn_start_epoch
+
+    def cvt_syncbn(self, runner):
+        if is_parallel(runner.model.module):
+            runner.model.module.module=\
+                SyncBatchNorm.convert_sync_batchnorm(runner.model.module.module,
+                                                     process_group=None)
+        else:
+            runner.model.module=\
+                SyncBatchNorm.convert_sync_batchnorm(runner.model.module,
+                                                     process_group=None)
+
+    def before_train_epoch(self, runner):
+        if runner.epoch>= self.syncbn_start_epoch and not self.is_syncbn:
+            print('start use syncbn')
+            self.cvt_syncbn(runner)
+            self.is_syncbn=True
+

docker-hub/FlashOCC/Flashocc/projects/mmdet3d_plugin/core/hook/utils.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from torch import nn
+
+__all__ = ['is_parallel']
+
+
+def is_parallel(model):
+    """check if model is in parallel mode."""
+    parallel_type = (
+        nn.parallel.DataParallel,
+        nn.parallel.DistributedDataParallel,
+    )
+    return isinstance(model, parallel_type)