Linhz/ViMNer
1
1# coding=utf-8
2# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16"""PyTorch BERT model."""
17
18from __future__ import absolute_import, division, print_function, unicode_literals
19
20import copy
21import json
22import logging
23import math
24import os
25import shutil
26import tarfile
27import tempfile
28import sys
29from io import open
30from torchcrf import CRF
31
32import torch
33from torch import nn
34from torch.nn import CrossEntropyLoss
35
36import torch.nn.functional as F
37from torch.autograd import Variable
38
39logger = logging.getLogger(__name__)
40
41
42def gelu(x):
43 """Implementation of the gelu activation function.
44 For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
45 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
46 Also see https://arxiv.org/abs/1606.08415
47 """
48 return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
49
50
51def swish(x):
52 return x * torch.sigmoid(x)
53
54
55ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish}
56
57from transformers import RobertaModel
58from transformers.models.roberta.modeling_roberta import RobertaLayer, RobertaPreTrainedModel, RobertaOutput, \
59 RobertaSelfOutput, RobertaIntermediate
60
61
62class RobertaSelfEncoder(nn.Module):
63 def __init__(self, config):
64 super(RobertaSelfEncoder, self).__init__()
65 layer = RobertaLayer(config)
66 self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(1)])
67
68 def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True):
69 all_encoder_layers = []
70 for layer_module in self.layer:
71 hidden_states = layer_module(hidden_states, attention_mask)
72 if output_all_encoded_layers:
73 all_encoder_layers.append(hidden_states)
74 if not output_all_encoded_layers:
75 all_encoder_layers.append(hidden_states)
76 return all_encoder_layers
77
78
79class RobertaCrossEncoder(nn.Module):
80 def __init__(self, config, layer_num):
81 super(RobertaCrossEncoder, self).__init__()
82 layer = RobertaCrossAttentionLayer(config)
83 self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(layer_num)])
84
85 def forward(self, s1_hidden_states, s2_hidden_states, s2_attention_mask, output_all_encoded_layers=True):
86 all_encoder_layers = []
87 for layer_module in self.layer:
88 s1_hidden_states = layer_module(s1_hidden_states, s2_hidden_states, s2_attention_mask)
89 if output_all_encoded_layers:
90 all_encoder_layers.append(s1_hidden_states)
91 if not output_all_encoded_layers:
92 all_encoder_layers.append(s1_hidden_states)
93 return all_encoder_layers
94
95
96class RobertaCoAttention(nn.Module):
97 def __init__(self, config):
98 super(RobertaCoAttention, self).__init__()
99 if config.hidden_size % config.num_attention_heads != 0:
100 raise ValueError(
101 "The hidden size (%d) is not a multiple of the number of attention "
102 "heads (%d)" % (config.hidden_size, config.num_attention_heads))
103 self.num_attention_heads = config.num_attention_heads
104 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
105 self.all_head_size = self.num_attention_heads * self.attention_head_size
106
107 self.query = nn.Linear(config.hidden_size, self.all_head_size)
108 self.key = nn.Linear(config.hidden_size, self.all_head_size)
109 self.value = nn.Linear(config.hidden_size, self.all_head_size)
110
111 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
112
113 def transpose_for_scores(self, x):
114 new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
115 x = x.view(*new_x_shape)
116 return x.permute(0, 2, 1, 3)
117
118 def forward(self, s1_hidden_states, s2_hidden_states, s2_attention_mask):
119 mixed_query_layer = self.query(s1_hidden_states)
120 mixed_key_layer = self.key(s2_hidden_states)
121 mixed_value_layer = self.value(s2_hidden_states)
122
123 query_layer = self.transpose_for_scores(mixed_query_layer)
124 key_layer = self.transpose_for_scores(mixed_key_layer)
125 value_layer = self.transpose_for_scores(mixed_value_layer)
126
127 # Take the dot product between "query" and "key" to get the raw attention scores.
128 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
129
130 attention_scores = attention_scores / math.sqrt(self.attention_head_size)
131 # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
132 attention_scores = attention_scores + s2_attention_mask
133
134 # Normalize the attention scores to probabilities.
135 attention_probs = nn.Softmax(dim=-1)(attention_scores)
136
137 # This is actually dropping out entire tokens to attend to, which might
138 # seem a bit unusual, but is taken from the original Transformer paper.
139 attention_probs = self.dropout(attention_probs)
140
141 context_layer = torch.matmul(attention_probs, value_layer)
142
143 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
144
145 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
146 context_layer = context_layer.view(*new_context_layer_shape)
147 return context_layer
148
149
150class RobertaCrossAttention(nn.Module):
151 def __init__(self, config):
152 super(RobertaCrossAttention, self).__init__()
153 self.self = RobertaCoAttention(config)
154 self.output = RobertaSelfOutput(config)
155
156 def forward(self, s1_input_tensor, s2_input_tensor, s2_attention_mask):
157 s1_cross_output = self.self(s1_input_tensor, s2_input_tensor, s2_attention_mask)
158 attention_output = self.output(s1_cross_output, s1_input_tensor)
159 return attention_output
160
161
162class RobertaCrossAttentionLayer(nn.Module):
163 def __init__(self, config):
164 super(RobertaCrossAttentionLayer, self).__init__()
165 self.attention = RobertaCrossAttention(config)
166 self.intermediate = RobertaIntermediate(config)
167 self.output = RobertaOutput(config)
168
169 def forward(self, s1_hidden_states, s2_hidden_states, s2_attention_mask):
170 attention_output = self.attention(s1_hidden_states, s2_hidden_states, s2_attention_mask)
171 intermediate_output = self.intermediate(attention_output)
172 layer_output = self.output(intermediate_output, attention_output)
173 return layer_output
174
175
176class UMT(RobertaPreTrainedModel):
177 """Coupled Cross-Modal Attention BERT model for token-level classification with CRF on top.
178 """
179
180 def __init__(self, config, layer_num1=1, layer_num2=1, layer_num3=1, num_labels_=2, auxnum_labels=2):
181 super(UMT, self).__init__(config)
182 self.num_labels = num_labels_
183 self.roberta = RobertaModel(config)
184 # self.trans_matrix = torch.zeros(num_labels, auxnum_labels)
185 self.self_attention = RobertaSelfEncoder(config)
186 self.self_attention_v2 = RobertaSelfEncoder(config)
187 self.dropout = nn.Dropout(config.hidden_dropout_prob)
188 self.vismap2text = nn.Linear(2048, config.hidden_size)
189 self.vismap2text_v2 = nn.Linear(2048, config.hidden_size)
190 self.txt2img_attention = RobertaCrossEncoder(config, layer_num1)
191 self.img2txt_attention = RobertaCrossEncoder(config, layer_num2)
192 self.txt2txt_attention = RobertaCrossEncoder(config, layer_num3)
193 self.gate = nn.Linear(config.hidden_size * 2, config.hidden_size)
194 ### self.self_attention = BertLastSelfAttention(config)
195 self.classifier = nn.Linear(config.hidden_size * 2, num_labels_)
196 self.aux_classifier = nn.Linear(config.hidden_size, auxnum_labels)
197
198 self.crf = CRF(num_labels_, batch_first=True)
199 self.aux_crf = CRF(auxnum_labels, batch_first=True)
200
201 self.init_weights()
202
203 # this forward is just for predict, not for train
204 # dont confuse this with _forward_alg above.
205 def forward(self, input_ids, segment_ids, input_mask, added_attention_mask, visual_embeds_att, trans_matrix,
206 labels=None, auxlabels=None):
207 # Get the emission scores from the BiLSTM
208 features = self.roberta(input_ids, token_type_ids=segment_ids,
209 attention_mask=input_mask) # batch_size * seq_len * hidden_size
210 sequence_output = features["last_hidden_state"]
211 sequence_output = self.dropout(sequence_output)
212
213 extended_txt_mask = input_mask.unsqueeze(1).unsqueeze(2)
214 extended_txt_mask = extended_txt_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
215 extended_txt_mask = (1.0 - extended_txt_mask) * -10000.0
216 aux_addon_sequence_encoder = self.self_attention(sequence_output, extended_txt_mask)
217
218 aux_addon_sequence_output = aux_addon_sequence_encoder[-1]
219 aux_addon_sequence_output = aux_addon_sequence_output[0]
220 aux_bert_feats = self.aux_classifier(aux_addon_sequence_output)
221 #######aux_bert_feats = self.aux_classifier(sequence_output)
222 trans_matrix_tensor = torch.tensor(trans_matrix, dtype=torch.float32, device=aux_bert_feats.device)
223 trans_bert_feats = torch.matmul(aux_bert_feats, trans_matrix_tensor)
224
225 # trans_bert_feats = torch.matmul(aux_bert_feats, trans_matrix.float())
226
227 main_addon_sequence_encoder = self.self_attention_v2(sequence_output, extended_txt_mask)
228 main_addon_sequence_output = main_addon_sequence_encoder[-1]
229 main_addon_sequence_output = main_addon_sequence_output[0]
230 vis_embed_map = visual_embeds_att.view(-1, 2048, 49).permute(0, 2, 1) # self.batch_size, 49, 2048
231 converted_vis_embed_map = self.vismap2text(vis_embed_map) # self.batch_size, 49, hidden_dim
232
233 # '''
234 # apply txt2img attention mechanism to obtain image-based text representations
235 img_mask = added_attention_mask[:, :49]
236 extended_img_mask = img_mask.unsqueeze(1).unsqueeze(2)
237 extended_img_mask = extended_img_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
238 extended_img_mask = (1.0 - extended_img_mask) * -10000.0
239
240 cross_encoder = self.txt2img_attention(main_addon_sequence_output, converted_vis_embed_map, extended_img_mask)
241 cross_output_layer = cross_encoder[-1] # self.batch_size * text_len * hidden_dim
242
243 # apply img2txt attention mechanism to obtain multimodal-based text representations
244 converted_vis_embed_map_v2 = self.vismap2text_v2(vis_embed_map) # self.batch_size, 49, hidden_dim
245
246 cross_txt_encoder = self.img2txt_attention(converted_vis_embed_map_v2, main_addon_sequence_output,
247 extended_txt_mask)
248 cross_txt_output_layer = cross_txt_encoder[-1] # self.batch_size * 49 * hidden_dim
249 cross_final_txt_encoder = self.txt2txt_attention(main_addon_sequence_output, cross_txt_output_layer,
250 extended_img_mask)
251 ##cross_final_txt_encoder = self.txt2txt_attention(aux_addon_sequence_output, cross_txt_output_layer, extended_img_mask)
252 cross_final_txt_layer = cross_final_txt_encoder[-1] # self.batch_size * text_len * hidden_dim
253 # cross_final_txt_layer = torch.add(cross_final_txt_layer, sequence_output)
254
255 # visual gate
256 merge_representation = torch.cat((cross_final_txt_layer, cross_output_layer), dim=-1)
257 gate_value = torch.sigmoid(self.gate(merge_representation)) # batch_size, text_len, hidden_dim
258 gated_converted_att_vis_embed = torch.mul(gate_value, cross_output_layer)
259 # reverse_gate_value = torch.neg(gate_value).add(1)
260 # gated_converted_att_vis_embed = torch.add(torch.mul(reverse_gate_value, cross_final_txt_layer),
261 # torch.mul(gate_value, cross_output_layer))
262
263 # direct concatenation
264 # gated_converted_att_vis_embed = self.dropout(gated_converted_att_vis_embed)
265 final_output = torch.cat((cross_final_txt_layer, gated_converted_att_vis_embed), dim=-1)
266 ###### final_output = self.dropout(final_output)
267 # middle_output = torch.cat((cross_final_txt_layer, gated_converted_att_vis_embed), dim=-1)
268 # final_output = torch.cat((sequence_output, middle_output), dim=-1)
269
270 ###### addon_sequence_output = self.self_attention(final_output, extended_txt_mask)
271 bert_feats = self.classifier(final_output)
272
273 alpha = 0.5
274 final_bert_feats = torch.add(torch.mul(bert_feats, alpha), torch.mul(trans_bert_feats, 1 - alpha))
275
276 # suggested by Hongjie
277 # bert_feats = F.log_softmax(bert_feats, dim=-1)
278
279 if labels is not None:
280 beta = 0.5 # 73.87(73.50) 85.37(85.00) 0.5 5e-5 #73.45 85.05 1.0 1 1 1 4e-5 # 73.63 0.1 1 1 1 5e-5 # old 0.1 2 1 1 85.23 0.2 1 1 85.04
281 ##beta = 0.6
282 aux_loss = - self.aux_crf(aux_bert_feats, auxlabels, mask=input_mask.byte(), reduction='mean')
283 main_loss = - self.crf(final_bert_feats, labels, mask=input_mask.byte(), reduction='mean')
284 loss = main_loss + beta * aux_loss
285 return loss
286 else:
287 pred_tags = self.crf.decode(final_bert_feats, mask=input_mask.byte())
288 return pred_tags
289
290
291