Method forward

openrec/modeling/decoders/bus_decoder.py:78–133

Args: tokens: (N, T, C) where T is length, N is batch size and C is classes number lengths: (N,)

(self, img_feat, data=None)

        return logits

    def forward(self, img_feat, data=None):
        """
        Args:
            tokens: (N, T, C) where T is length, N is batch size and C is classes number
            lengths: (N,)
        """
        img_feat = img_feat + self.v_embeding
        B, L, C = img_feat.shape

        # --------------------------------------------------------------------------
        # decoder procedure
        T = self.max_length
        zeros = img_feat.new_zeros((B, T, C))
        zeros_len = img_feat.new_zeros(B)
        query = self.pos_encoder(zeros)

        # 1. vision decode
        v_embed = torch.cat((img_feat, self.l_mask.repeat(B, T, 1)),
                            dim=1)  # v
        padding_mask = _get_mask(
            self.max_length + zeros_len,
            self.max_length)  # 对tokens长度以外的padding # B, maxlen maxlen
        v_mask = torch.zeros((1, 1, self.max_length, L),
                             device=img_feat.device).tile([B, 1, 1,
                                                           1])  # maxlen L
        mask = torch.cat((v_mask, padding_mask), 3)
        v_logits = self.forward_decoder(query, v_embed, mask=mask)

        # 2. language decode
        if self.training and self.pretraining:
            tgt = torch.where(data[0] == self.ignore_index, 0, data[0])
            tokens = F.one_hot(tgt, num_classes=self.out_channels)
            tokens = tokens.float()
            lengths = data[-1]
        else:
            tokens = torch.softmax(v_logits, dim=-1)
            lengths = _get_length(v_logits)
            tokens = tokens.detach()
        token_embed = self.proj(tokens)  # (N, T, E)
        token_embed = self.token_encoder(token_embed)  # (T, N, E)
        token_embed = token_embed + self.l_embeding

        padding_mask = _get_mask(lengths,
                                 self.max_length)  # 对tokens长度以外的padding
        mask = torch.cat((v_mask, padding_mask), 3)
        l_embed = torch.cat((self.v_mask.repeat(B, L, 1), token_embed), dim=1)
        l_logits = self.forward_decoder(query, l_embed, mask=mask)

        # 3. vision language decode
        vl_embed = torch.cat((img_feat, token_embed), dim=1)
        vl_logits = self.forward_decoder(query, vl_embed, mask=mask)

        if self.training:
            return {'align': [vl_logits], 'lang': l_logits, 'vision': v_logits}
        else:
            return F.softmax(vl_logits, -1)