jefferyyu / tombert Goto Github PK

class MBertForMMSequenceClassification(PreTrainedBertModel):
    """BERT model for classification with text and image inputs, pooling-1+text (MBERT I)
    """

def __init__(self, config, num_labels=2, pooling="cls"):
    super(MBertForMMSequenceClassification, self).__init__(config)
    self.num_labels = num_labels
    self.pooling = pooling
    self.bert = BertModel(config)
    self.dropout = nn.Dropout(config.hidden_dropout_prob)
    self.vismap2text = nn.Linear(2048, config.hidden_size)
    #self.img_attention = BertLayer(config)
    self.comb_attention = MultimodalEncoder(config)
    if pooling == "cls":
        self.text_pooler = BertText1Pooler(config)
        self.classifier = nn.Linear(config.hidden_size, num_labels)
    elif pooling == "first":
        self.img_pooler = BertPooler(config)
        self.classifier = nn.Linear(config.hidden_size, num_labels)
    else:
        self.text_pooler = BertText1Pooler(config)
        self.img_pooler = BertPooler(config)
        self.classifier = nn.Linear(config.hidden_size * 2, num_labels)

    self.apply(self.init_bert_weights)

def forward(self, input_ids, s2_input_ids, visual_embeds_att, token_type_ids=None, s2_type_ids=None,
            attention_mask=None, s2_mask=None, added_attention_mask=None, labels=None, copy_flag=False):
    # Concatenate Bert-based Text, Text-Aware Image and Image-Aware Text
    sequence_output, pooled_output = self.bert(input_ids, token_type_ids, attention_mask,
                                               output_all_encoded_layers=False)

    # apply entity-based attention mechanism to obtain different image representations
    vis_embed_map = visual_embeds_att.view(-1, 2048, 49).permute(0, 2, 1)  # self.batch_size, 49, 2048
    vis_pooled_output, _ = vis_embed_map.max(1)  # self.batch_size, 2048
    converted_vis_embed_map = self.vismap2text(vis_pooled_output)  # self.batch_size, hidden_dim

    transpose_img_embed = converted_vis_embed_map.unsqueeze(1)
    text_img_output = torch.cat((transpose_img_embed, sequence_output), dim=1)

    comb_attention_mask = added_attention_mask[:, 48:]  # only the first dimension is for image
    extended_attention_mask = comb_attention_mask.unsqueeze(1).unsqueeze(2)
    extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0