generic_classifier.py (13340B)
1 import numpy as np 2 import torch 3 import transformers 4 from sentence_transformers import SentenceTransformer 5 6 import mmsrl.loss 7 import mmsrl.classification 8 9 class Model(torch.nn.Module): 10 def __init__(self, config, label_weights, model=None): 11 super().__init__() 12 self.config = config 13 assert self.config.use_clip == True or self.config.use_visualbert == True 14 15 if self.config.features_path=="features_b7": 16 self.features_dim = 2560 17 elif self.config.features_path == "features_vgg": 18 self.features_dim = 25088 if self.config.pooling=='mlp' else 512 19 elif self.config.features_path=="all": 20 self.features_dim = 29696 21 elif self.config.features_path=="features_frcnn": 22 self.features_dim = 2048 23 24 self.model = model 25 if self.config.use_clip: output_emb_dim = self.model.ln_final.bias.shape[0] 26 if self.config.unfreeze == "classifier": 27 self.model.requires_grad_(False) 28 if self.config.use_visualbert: 29 self.image_linear = torch.nn.Linear(self.features_dim, 512) 30 self.transformer = transformers.VisualBertModel.from_pretrained(self.config.visualbert_name, add_pooling_layer = True) 31 if not self.config.use_clip: 32 output_emb_dim = self.transformer.config.hidden_size 33 if self.config.unfreeze == "classifier": 34 self.transformer.requires_grad_(False) 35 36 if self.config.features_path: 37 self.dense_vgg_1024 = torch.nn.Linear(self.features_dim, 1024) 38 self.dense_img_feat = torch.nn.Linear(1024, output_emb_dim) 39 self.drop20 = torch.nn.Dropout(p=0.2) 40 if self.config.use_sbert: 41 self.model_sent_trans = SentenceTransformer(self.config.sbert_model) 42 43 if self.config.pooling == "attention": 44 self.linear_key = torch.nn.Linear(output_emb_dim, self.config.hidden_size) 45 self.linear_value = torch.nn.Linear(output_emb_dim, self.config.hidden_size) 46 47 self.entities_linear = torch.nn.Linear(output_emb_dim, self.config.hidden_size) 48 self.classifier = torch.nn.Linear(self.config.hidden_size, 4) # output * number of classes 49 self.loss = mmsrl.loss.Loss(self.config, label_weights) 50 self.relu = torch.nn.ReLU() 51 52 self.num_features = 0 53 if self.config.use_clip: self.num_features +=2 54 if self.config.use_visualbert: self.num_features +=1 55 if self.config.features_path and self.config.use_clip: self.num_features +=1 56 if self.config.use_caption: self.num_features +=1 57 if self.config.use_entities: self.num_features +=1 58 print("Number of features in total: ", self.num_features) 59 60 if self.config.pooling == "attention": 61 self.segment_embeddings = torch.nn.Embedding(self.num_features, self.config.hidden_size) 62 elif self.config.pooling != "mlp": 63 raise RuntimeError(f"Unknown config value for config.pooling in clip: {self.config.pooling}") 64 65 #self.model.transformer.resblocks[-1].requires_grad_(True) 66 #self.model.visual.transformer.resblocks[-1].requires_grad_(True) 67 68 def train(self, value=True): 69 super().train(value) 70 self.model.eval() 71 if self.config.use_visualbert: 72 self.transformer.eval() 73 74 def clip_encode_text_without_pooling(self, text): 75 """ This code comes from clip, except that we removed the selection along text length. """ 76 x = self.model.token_embedding(text).type(self.model.dtype) # [batch_size, n_ctx, d_model] 77 78 x = x + self.model.positional_embedding.type(self.model.dtype) 79 x = x.permute(1, 0, 2) # NLD -> LND 80 x = self.model.transformer(x) 81 x = x.permute(1, 0, 2) # LND -> NLD 82 x = self.model.ln_final(x).type(self.model.dtype) 83 84 # x.shape = [batch_size, n_ctx, transformer.width] 85 x = x @ self.model.text_projection 86 return x 87 88 def clip_encode_image_without_pooling(self, image): 89 """ This code comes from clip, except that we removed the selection along image length. """ 90 x = image.to(dtype=self.model.dtype) 91 x = self.model.visual.conv1(x) # shape = [*, width, grid, grid] 92 x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2] 93 x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width] 94 x = torch.cat([self.model.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width] 95 x = x + self.model.visual.positional_embedding.to(x.dtype) 96 x = self.model.visual.ln_pre(x) 97 98 x = x.permute(1, 0, 2) # NLD -> LND 99 x = self.model.visual.transformer(x) 100 x = x.permute(1, 0, 2) # LND -> NLD 101 102 x = self.model.visual.ln_post(x) 103 104 if self.model.visual.proj is not None: 105 x = x @ self.model.visual.proj 106 107 return x 108 109 def forward(self, input_ids, image, attention_mask, entities_mask, entities_input_ids, entities_attention_mask, entities_name, image_feature = None, vbert_attention_mask = None, vbert_input_ids = None, caption_input_ids = None, caption_attention_mask=None, labels=None, **kwargs): 110 # entities are tri-dimensionnal (batch * nb of entities * emb dim) while text is bi-dimensional 111 nb_entities = entities_input_ids.shape[1] 112 batch_size = entities_input_ids.shape[0] 113 fake_batch_size = batch_size*nb_entities 114 115 features = [] 116 masks =[] 117 segments = [] 118 119 # 1) Encoding from CLIP 120 if self.config.use_clip: 121 num_feat = len(features) 122 if self.config.pooling == "mlp": 123 clip_text_features = self.model.encode_text(input_ids) 124 clip_image_features = self.model.encode_image(image) 125 elif self.config.pooling == "attention": 126 clip_text_features = self.clip_encode_text_without_pooling(input_ids) 127 clip_image_features = self.clip_encode_image_without_pooling(image) 128 129 features.append(clip_text_features) 130 features.append(clip_image_features) 131 attention_mask = input_ids!=0 132 segments.append(attention_mask.new_ones(batch_size, clip_text_features.shape[1], dtype=torch.int64)*num_feat) 133 segments.append(attention_mask.new_ones(batch_size, clip_image_features.shape[1], dtype=torch.int64)*num_feat) 134 masks.append(attention_mask) 135 masks.append(attention_mask.new_ones(batch_size, clip_image_features.shape[1])) 136 137 # 2) Image features pre-extracted using various models 138 if self.config.features_path and self.config.use_clip: # visualbert already use them 139 num_feat = len(features) 140 image_feature_tensor = torch.tensor(np.array(image_feature), device=input_ids.device) 141 image_feature_tensor = self.drop20(torch.nn.functional.relu(self.dense_img_feat(self.drop20(torch.nn.functional.relu(self.dense_vgg_1024(image_feature_tensor)))))) 142 if self.config.pooling == "mlp": 143 image_feature_tensor = image_feature_tensor.squeeze(1) 144 features.append(image_feature_tensor) 145 # All image features have the same shape 146 segments.append(attention_mask.new_ones(batch_size, image_feature_tensor.shape[1], dtype=torch.int64)*num_feat) 147 masks.append(attention_mask.new_ones(batch_size, image_feature_tensor.shape[1])) 148 149 # 3) List of all entities of the meme, with one embedding per entity, concatenated. 150 if self.config.use_entities: 151 num_feat = len(features) 152 if self.config.get("use_sbert"): 153 ent = torch.vstack([torch.mean(torch.tensor(self.model_sent_trans.encode(entities)), dim=0) for entities in entities_name]) 154 entities_list_embedding = torch.stack(tuple([torch.tensor(self.model_sent_trans.encode(entities + ['']*(entities_mask.shape[1] - len(entities)))) for entities in entities_name])).to(device = input_ids.device) 155 else: 156 entities_list_embedding = self.model.encode_text(entities_input_ids.view(fake_batch_size, -1)).view(batch_size, nb_entities, -1) 157 #entities_feat = self.drop5(torch.nn.functional.relu(entities_embedding)) # batch * nb_ent * emb_dim 158 if self.config.pooling == "mlp": 159 entities_list_embedding = (entities_list_embedding*entities_mask.unsqueeze(2)).sum(1)/entities_mask.sum(1,keepdim=True) 160 features.append(entities_list_embedding) 161 segments.append(entities_mask.new_ones(batch_size, entities_list_embedding.shape[1], dtype=torch.int64)*num_feat) 162 masks.append(entities_mask) 163 164 # 4) Caption of the meme 165 if self.config.use_caption: 166 num_feat = len(features) 167 # should we use sentenceBERT and just concat it with the clip_text_emb ? 168 #cap = torch.tensor(self.model_sent_trans.encode(caption)).to(device = input_ids.device) 169 #text_embedding_pooled = self.caption_entiy_linear(torch.cat((entities_embedding_pooled, self.caption_linear(cap)), dim=1)) 170 if self.config.pooling == "mlp": 171 caption_embedding = self.model.encode_text(caption_input_ids) 172 elif self.config.pooling == "attention": 173 caption_embedding = self.clip_encode_text_without_pooling(caption_input_ids) 174 features.append(caption_embedding) 175 caption_attention_mask = caption_input_ids!=0 176 segments.append(caption_attention_mask.new_ones(batch_size, caption_embedding.shape[1], dtype=torch.int64)*num_feat) 177 masks.append(caption_attention_mask) 178 179 # 5) Features from visualBERT: 180 if self.config.use_visualbert: 181 num_feat = len(features) 182 visual_embeds = self.image_linear(torch.tensor(np.array(image_feature), device=input_ids.device).to(device=input_ids.device)) 183 visual_attention_mask = torch.ones(visual_embeds.shape[:-1], dtype=torch.long).to(device=input_ids.device) 184 visual_token_type_ids = torch.ones(visual_embeds.shape[:-1], dtype=torch.long).to(device=input_ids.device) 185 token_type_ids = torch.ones(vbert_attention_mask.shape, dtype=torch.long).to(device=input_ids.device) 186 outputs = self.transformer(input_ids=vbert_input_ids, attention_mask=vbert_attention_mask, token_type_ids=token_type_ids, 187 visual_embeds=visual_embeds, visual_attention_mask=visual_attention_mask, visual_token_type_ids=visual_token_type_ids, output_hidden_states=True, return_dict=True) 188 if self.config.pooling == "mlp": 189 features.append(outputs.pooler_output) 190 elif self.config.pooling == "attention": 191 features.append(outputs.last_hidden_state) 192 segments.append(outputs.last_hidden_state.new_ones(batch_size, outputs.last_hidden_state.shape[1], dtype=torch.int64)*num_feat) 193 masks.append(entities_mask.new_ones(batch_size, outputs.last_hidden_state.shape[1])) 194 195 # 6) Entity feature to be used as query 196 entities_features = self.model.encode_text(entities_input_ids.view(fake_batch_size, -1)).view(batch_size, nb_entities, -1) 197 entities_embedding = self.entities_linear(entities_features.to(dtype=self.entities_linear.weight.dtype)) 198 199 assert self.num_features == len(features) 200 features_value = list(map(self.linear_value, [f.to(dtype=self.linear_value.weight.dtype) for f in features])) 201 202 # 7) MLP and attention 203 if self.config.pooling == "mlp": 204 features_value = sum(features_value) 205 # We want the joint representation of the features and entities. 206 hidden = self.relu(features_value.unsqueeze(1) + entities_embedding / len(features) +1 ) 207 elif self.config.pooling == "attention": 208 features_key = list(map(self.linear_key, [f.to(dtype=self.linear_key.weight.dtype) for f in features])) 209 210 memory_key = torch.cat(features_key, dim=1) + self.segment_embeddings(torch.cat(segments, dim=1)) 211 memory_value = torch.cat(features_value, dim=1) + self.segment_embeddings(torch.cat(segments, dim=1)) 212 mask = torch.cat(masks, dim=1) 213 214 attention = torch.einsum("blf,bef->bel", memory_key, entities_embedding) 215 attention_entities_mask = entities_mask.unsqueeze(2) & mask.unsqueeze(1) # make the softmax only on entities and words --> be->be1 & bl->b1l --> bel 216 attention_max = torch.max(attention, dim=2, keepdim=True)[0] # max along sentence length 217 attention_exp = torch.exp(attention-attention_max) 218 attention_exp = attention_exp * attention_entities_mask.float() 219 attention_softmax = attention_exp / (attention_exp.sum(dim=2,keepdim=True) + (1-attention_entities_mask.float())) 220 221 hidden = torch.sum(attention_softmax.unsqueeze(3) * memory_value.unsqueeze(1), dim=2) # sum along sentence length: torch.Size([8, nbent, seqlen, 1]) + torch.Size([8, 1, seqlen, f] --> torch.Size([8, nbent, f]) 222 hidden = self.relu(hidden) 223 224 logits = self.classifier(hidden) # batch * entity * nb_labels (4) 225 return self.loss(logits, labels)