def _init_layers(self):
        """Initialize classification branch and regression branch of head."""
        cls_branch = []
        for _ in range(self.num_reg_fcs):
            cls_branch.append(Linear(self.embed_dims, self.embed_dims))
            cls_branch.append(nn.LayerNorm(self.embed_dims))
            cls_branch.append(nn.ReLU(inplace=True))
        cls_branch.append(Linear(self.embed_dims, self.cls_out_channels))
        cls_branch = nn.Sequential(*cls_branch)

        reg_branch = []
        for _ in range(self.num_reg_fcs):
            reg_branch.append(Linear(self.embed_dims, self.embed_dims))
            reg_branch.append(nn.ReLU())
        reg_branch.append(Linear(self.embed_dims, self.code_size))
        reg_branch = nn.Sequential(*reg_branch)

        mot_reg_branch = []
        for _ in range(self.num_reg_fcs):
            mot_reg_branch.append(Linear(self.embed_dims*2, self.embed_dims*2))
            mot_reg_branch.append(nn.ReLU())
        mot_reg_branch.append(Linear(self.embed_dims*2, self.fut_ts*2))
        mot_reg_branch = nn.Sequential(*mot_reg_branch)

        mot_cls_branch = []
        for _ in range(self.num_reg_fcs):
            mot_cls_branch.append(Linear(self.embed_dims*2, self.embed_dims*2))
            mot_cls_branch.append(nn.LayerNorm(self.embed_dims*2))
            mot_cls_branch.append(nn.ReLU(inplace=True))
        mot_cls_branch.append(Linear(self.embed_dims*2, self.mot_num_cls))
        mot_cls_branch = nn.Sequential(*mot_cls_branch)

        map_cls_branch = []
        for _ in range(self.num_reg_fcs):
            map_cls_branch.append(Linear(self.embed_dims, self.embed_dims))
            map_cls_branch.append(nn.LayerNorm(self.embed_dims))
            map_cls_branch.append(nn.ReLU(inplace=True))
        map_cls_branch.append(Linear(self.embed_dims, self.map_cls_out_channels))
        map_cls_branch = nn.Sequential(*map_cls_branch)

        map_reg_branch = []
        for _ in range(self.num_reg_fcs):
            map_reg_branch.append(Linear(self.embed_dims, self.embed_dims))
            map_reg_branch.append(nn.ReLU())
        map_reg_branch.append(Linear(self.embed_dims, self.map_code_size))
        map_reg_branch = nn.Sequential(*map_reg_branch)


        ego_query_pre_branch = []
        ego_query_pre_branch.append(Linear(self.embed_dims * self.fut_ts, self.embed_dims))
        ego_query_pre_branch.append(nn.ReLU())
        ego_query_pre_branch.append(Linear(self.embed_dims, self.embed_dims))
        self.ego_query_pre_branch = nn.Sequential(*ego_query_pre_branch)


        def _get_clones(module, N):
            return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

        # last reg_branch is used to generate proposal from
        # encode feature map when as_two_stage is True.
        num_decoder_layers = 1
        num_map_decoder_layers = 1
        if self.transformer.decoder is not None:
            num_decoder_layers = self.transformer.decoder.num_layers
        if self.transformer.map_decoder is not None:
            num_map_decoder_layers = self.transformer.map_decoder.num_layers
        num_mot_decoder_layers = 1
        num_pred = (num_decoder_layers + 1) if \
            self.as_two_stage else num_decoder_layers
        mot_num_pred = (num_mot_decoder_layers + 1) if \
            self.as_two_stage else num_mot_decoder_layers
        map_num_pred = (num_map_decoder_layers + 1) if \
            self.as_two_stage else num_map_decoder_layers

        if self.with_box_refine:
            self.cls_branches = _get_clones(cls_branch, num_pred)
            self.reg_branches = _get_clones(reg_branch, num_pred)
            self.mot_reg_branches = _get_clones(mot_reg_branch, mot_num_pred)
            self.mot_cls_branches = _get_clones(mot_cls_branch, mot_num_pred)
            self.map_cls_branches = _get_clones(map_cls_branch, map_num_pred)
            self.map_reg_branches = _get_clones(map_reg_branch, map_num_pred)
        else:
            self.cls_branches = nn.ModuleList(
                [cls_branch for _ in range(num_pred)])
            self.reg_branches = nn.ModuleList(
                [reg_branch for _ in range(num_pred)])
            self.mot_reg_branches = nn.ModuleList(
                [mot_reg_branch for _ in range(mot_num_pred)])
            self.mot_cls_branches = nn.ModuleList(
                [mot_cls_branch for _ in range(mot_num_pred)])
            self.map_cls_branches = nn.ModuleList(
                [map_cls_branch for _ in range(map_num_pred)])
            self.map_reg_branches = nn.ModuleList(
                [map_reg_branch for _ in range(map_num_pred)])

        if not self.as_two_stage:
            self.bev_embedding = nn.Embedding(
                self.bev_h * self.bev_w, self.embed_dims)
            self.query_embedding = nn.Embedding(self.num_query,
                                                self.embed_dims * 2)
            if self.map_query_embed_type == 'all_pts':
                self.map_query_embedding = nn.Embedding(self.map_num_query,
                                                    self.embed_dims * 2)
            elif self.map_query_embed_type == 'instance_pts':
                self.map_query_embedding = None
                self.map_instance_embedding = nn.Embedding(self.map_num_vec, self.embed_dims * 2)
                self.map_pts_embedding = nn.Embedding(self.map_num_pts_per_vec, self.embed_dims * 2)
        
        if self.mot_decoder is not None:
            self.mot_decoder = build_transformer_layer_sequence(self.mot_decoder)
            self.mot_mode_query = nn.Embedding(self.mot_fut_mode, self.embed_dims)	
            self.mot_mode_query.weight.requires_grad = True
        else:
            raise NotImplementedError('Not implement yet')

        if self.mot_map_decoder is not None:
            self.lane_encoder = LaneNet(self.embed_dims, self.embed_dims // 2, 3)
            self.mot_map_decoder = build_transformer_layer_sequence(self.mot_map_decoder)

        # self.ego_query = nn.Embedding(self.plan_fut_mode, self.embed_dims)
        # self.ego_query = pos2posemb2d(self.plan_anchors.reshape(1, self.plan_fut_mode * self.fut_ts, -1)) \
        #         .reshape(self.plan_fut_mode, self.fut_ts, -1)

        if self.ego_pv_decoder is not None:
            self.ego_pv_decoder = build_transformer_layer_sequence(self.ego_pv_decoder)
            MAXNUM_PV_TOKEN = 800
            self.pv_pos_embedding = nn.Embedding(
                MAXNUM_PV_TOKEN, self.embed_dims)

        if self.ego_agent_decoder is not None:
            self.ego_agent_decoder = build_transformer_layer_sequence(self.ego_agent_decoder)

        if self.ego_map_decoder is not None:
            self.ego_map_decoder = build_transformer_layer_sequence(self.ego_map_decoder)

        plan_reg_branch = []
        # plan_in_dim = self.embed_dims*4 + len(self.ego_lcf_feat_idx) \
        #     if self.ego_lcf_feat_idx is not None else self.embed_dims*4
        for _ in range(self.num_reg_fcs):
            plan_reg_branch.append(Linear(self.embed_dims, self.embed_dims))
            plan_reg_branch.append(nn.ReLU())
        plan_reg_branch.append(Linear(self.embed_dims, self.fut_ts*2))
        self.plan_reg_branch = nn.Sequential(*plan_reg_branch)

        self.fus_mlp = nn.Sequential(
            nn.Linear(self.mot_fut_mode*2*self.embed_dims, self.embed_dims, bias=True),
            nn.LayerNorm(self.embed_dims),
            nn.ReLU(),
            nn.Linear(self.embed_dims, self.embed_dims, bias=True))

        if self.interaction_pe_type == 'sine_mlp':
            pe_embed_mlps = nn.Sequential(
                nn.Linear(self.embed_dims, self.embed_dims*2),
                nn.ReLU(),
                nn.Linear(self.embed_dims*2, self.embed_dims),
            )
        elif self.interaction_pe_type == 'mlp':
            pe_embed_mlps = nn.Linear(2, self.embed_dims)
        else:
            raise NotImplementedError('Not implement yet')
        
        self.pe_embed_mlps = _get_clones(pe_embed_mlps, 4)

        self.ego_feat_projs = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for agent
                nn.Sequential(
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for map
                nn.Sequential(
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for traffic light
                nn.Sequential(
                    nn.Linear(2, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for target point
                nn.Sequential(
                    nn.Linear(len(self.ego_lcf_feat_idx), self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for lcf feat
                nn.Sequential(
                    nn.Linear(6, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ),  # for cmdid
                nn.Sequential(
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for pv_feat
                nn.Sequential(
                    nn.Linear(140 * 140, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(self.embed_dims, self.embed_dims, bias=True),
                ), # for target point rasterized
            ]
        )

        # NOTE: add front-view feature encoder
        if self.cf_backbone is not None:
            self.cf_backbone = build_backbone(self.cf_backbone)

        tl_feats_branch = []
        tl_feats_branch.append(Linear(2048 * 10 * 13 + 6 * self.embed_dims * 5 * 7, self.embed_dims)) #10 * 13  5 * 7
        tl_feats_branch.append(nn.LayerNorm(self.embed_dims))
        tl_feats_branch.append(nn.ReLU(inplace=True))
        tl_feats_branch.append(Linear(self.embed_dims, self.embed_dims))
        tl_feats_branch.append(nn.LayerNorm(self.embed_dims))
        tl_feats_branch.append(nn.ReLU(inplace=True))
        self.tl_feats_branch = nn.Sequential(*tl_feats_branch)

        tl_status_cls_branch = []
        tl_status_cls_branch.append(Linear(self.embed_dims, self.embed_dims))
        tl_status_cls_branch.append(nn.LayerNorm(self.embed_dims))
        tl_status_cls_branch.append(nn.ReLU(inplace=True))
        tl_status_cls_branch.append(Linear(self.embed_dims, self.tl_status_num_cls))
        self.tl_status_cls_branch = nn.Sequential(*tl_status_cls_branch)

        tl_trigger_cls_branch = []
        tl_trigger_cls_branch.append(Linear(self.embed_dims, self.embed_dims))
        tl_trigger_cls_branch.append(nn.LayerNorm(self.embed_dims))
        tl_trigger_cls_branch.append(nn.ReLU(inplace=True))
        tl_trigger_cls_branch.append(Linear(self.embed_dims, self.tl_trigger_num_cls))
        self.tl_trigger_cls_branch = nn.Sequential(*tl_trigger_cls_branch)
        
        stopsign_trigger_cls_branch = []
        stopsign_trigger_cls_branch.append(Linear(self.embed_dims, self.embed_dims))
        stopsign_trigger_cls_branch.append(nn.LayerNorm(self.embed_dims))
        stopsign_trigger_cls_branch.append(nn.ReLU(inplace=True))
        stopsign_trigger_cls_branch.append(Linear(self.embed_dims, self.stopsign_trigger_num_cls))
        self.stopsign_trigger_cls_branch = nn.Sequential(*stopsign_trigger_cls_branch)

        plan_cls_col_branch = []
        # plan_cls_col_branch.append(Linear(self.embed_dims, self.embed_dims))
        # plan_cls_col_branch.append(nn.LayerNorm(self.embed_dims))
        # plan_cls_col_branch.append(nn.ReLU(inplace=True))
        plan_cls_col_branch.append(Linear(self.embed_dims, 1))
        self.plan_cls_col_branch = nn.Sequential(*plan_cls_col_branch)

        plan_cls_bd_branch = []
        # plan_cls_bd_branch.append(Linear(self.embed_dims, self.embed_dims))
        # plan_cls_bd_branch.append(nn.LayerNorm(self.embed_dims))
        # plan_cls_bd_branch.append(nn.ReLU(inplace=True))
        plan_cls_bd_branch.append(Linear(self.embed_dims, 1))
        self.plan_cls_bd_branch = nn.Sequential(*plan_cls_bd_branch)

        plan_cls_cl_branch = []
        # plan_cls_cl_branch.append(Linear(self.embed_dims, self.embed_dims))
        # plan_cls_cl_branch.append(nn.LayerNorm(self.embed_dims))
        # plan_cls_cl_branch.append(nn.ReLU(inplace=True))
        plan_cls_cl_branch.append(Linear(self.embed_dims, 1))
        self.plan_cls_cl_branch = nn.Sequential(*plan_cls_cl_branch)

        plan_cls_expert_branch = []
        plan_cls_expert_branch.append(Linear(self.embed_dims, self.embed_dims))
        plan_cls_expert_branch.append(nn.LayerNorm(self.embed_dims))
        plan_cls_expert_branch.append(nn.ReLU(inplace=True))
        plan_cls_expert_branch.append(Linear(self.embed_dims, 1))
        self.plan_cls_expert_branch = nn.Sequential(*plan_cls_expert_branch)
Initialize classification branch and regression branch of head.
VADv2/VADv2_head.py