Method forward

models/aios/transformer.py:266–498

(self,
                srcs,
                masks,
                refpoint_embed,
                pos_embeds,
                tgt,
                attn_mask=None,
                attn_mask2=None,
                attn_mask3=None)

    # srcs: features; refpoint_embed:
    def forward(self,
                srcs,
                masks,
                refpoint_embed,
                pos_embeds,
                tgt,
                attn_mask=None,
                attn_mask2=None,
                attn_mask3=None):
        # pdb.set_trace()
        # prepare input for encoder
        src_flatten = []
        mask_flatten = []
        lvl_pos_embed_flatten = []
        spatial_shapes = []
        for lvl, (src, mask, pos_embed) in enumerate(
                zip(srcs, masks, pos_embeds)):  # for feature level
            bs, c, h, w = src.shape
            spatial_shape = (h, w)
            spatial_shapes.append(spatial_shape)

            src = src.flatten(2).transpose(1, 2)  # bs, hw, c
            mask = mask.flatten(1)  # bs, hw
            pos_embed = pos_embed.flatten(2).transpose(1, 2)  # bs, hw, c
            if self.num_feature_levels > 1 and self.level_embed is not None:
                lvl_pos_embed = pos_embed + self.level_embed[lvl].view(
                    1, 1, -1)  # level_embed[lvl]: [256]
            else:
                lvl_pos_embed = pos_embed
            lvl_pos_embed_flatten.append(lvl_pos_embed)
            src_flatten.append(src)
            mask_flatten.append(mask)
        src_flatten = torch.cat(src_flatten, 1)  # bs, \sum{hxw}, c
        mask_flatten = torch.cat(mask_flatten, 1)  # bs, \sum{hxw}
        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten,
                                          1)  # bs, \sum{hxw}, c
        spatial_shapes = torch.as_tensor(spatial_shapes,
                                         dtype=torch.long,
                                         device=src_flatten.device)
        level_start_index = torch.cat((spatial_shapes.new_zeros(
            (1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
        valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
        # two stage
        if self.two_stage_type in ['early', 'combine']:
            output_memory, output_proposals = gen_encoder_output_proposals(
                src_flatten, mask_flatten, spatial_shapes)
            output_memory = self.enc_output_norm_backbone(
                self.enc_output_backbone(output_memory))

            # gather boxes
            topk = self.num_queries
            enc_outputs_class = self.encoder.class_embed[0](output_memory)
            enc_topk_proposals = torch.topk(enc_outputs_class.max(-1)[0],
                                            topk,
                                            dim=1)[1]  # bs, nq
            enc_refpoint_embed = torch.gather(
                output_proposals, 1,
                enc_topk_proposals.unsqueeze(-1).repeat(1, 1, 4))