ofa.py (9006B)
1 import math 2 3 import torch 4 5 import mmsrl.loss 6 import mmsrl.utils 7 8 9 class Model(torch.nn.Module): 10 def __init__(self, config, label_weights, ofa_task, ofa_config, ofa_models): 11 super().__init__() 12 self.config = config 13 self.ofa_task = ofa_task 14 self.ofa_config = ofa_config 15 assert(len(ofa_models) == 1) # Check if it's really an ensemble or a single model 16 self.model = ofa_models[0] 17 18 if self.config.get("ofa_classification_head"): 19 self.model.register_classification_head("mmsrl", 4) 20 self.loss = mmsrl.loss.Loss(self.config, label_weights) 21 else: 22 self.loss = torch.nn.NLLLoss(reduction="none") 23 24 if self.config.unfreeze != "all": 25 self.model.requires_grad_(False) 26 27 if self.config.unfreeze == "classifier": 28 if self.config.get("ofa_classification_head"): 29 self.model.classification_heads["mmsrl"].requires_grad_(True) 30 else: 31 self.model.decoder.output_projection.requires_grad_(True) 32 33 if self.config.get("freeze_embeddings"): 34 self.model.encoder.embed_tokens.requires_grad_(False) 35 self.model.decoder.embed_tokens.requires_grad_(False) 36 37 def train(self, mode=True): 38 super().train(mode) 39 if self.config.get("freeze_resnet"): 40 for module in self.model.encoder.embed_images.modules(): 41 if isinstance(module, torch.nn.BatchNorm2d): 42 module.eval() 43 module.requires_grad_(False) 44 45 def masked_log_softmax(self, logits, mask): 46 if mask is None: 47 return torch.nn.functional.log_softmax(logits, dim=2) 48 49 fake_batch_size: int = logits.shape[0] 50 sequence_length: int = logits.shape[1] 51 vocabulary_size: int = logits.shape[2] 52 mask = mask.view(fake_batch_size, sequence_length, vocabulary_size) 53 54 logits_max = torch.max(logits, dim=2, keepdim=True)[0] # max along vocabulary 55 logits_exp = torch.exp(logits - logits_max) 56 logits_exp = mask.float() * logits_exp 57 logits_sum = torch.sum(logits_exp, dim=2, keepdim=True) 58 logits_sum += (logits_sum==0).float() # avoid log(0) when everything is masked 59 log_softmax = logits - logits_max - torch.log(logits_sum) 60 return mask.float() * log_softmax 61 62 def forward(self, input_ids, length, image, image_mask, decoder_input_ids, decoder_target, decoder_target_attention_mask, labels=None, **kwargs): 63 if self.config.get("ofa_classification_head"): 64 if self.config.get("group_image_entities"): 65 batch_size: int = input_ids.shape[0] 66 num_entities: int = input_ids.shape[1] 67 fake_batch_size: int = batch_size * num_entities 68 input_ids = input_ids.view(fake_batch_size, -1) 69 length = length.view(fake_batch_size) 70 decoder_input_ids = decoder_input_ids.view(fake_batch_size, -1) 71 image = image.repeat_interleave(repeats=num_entities, dim=0) 72 image_mask = image_mask.repeat_interleave(repeats=num_entities, dim=0) 73 else: 74 input_ids = input_ids.squeeze(1) 75 length = length.squeeze(1) 76 decoder_input_ids = decoder_input_ids.squeeze(1) 77 78 outputs = self.model( 79 input_ids, 80 length, 81 decoder_input_ids, 82 patch_images=image, 83 patch_masks=image_mask[:, 0, 0], 84 classification_head_name="mmsrl")[0] 85 if input_ids.shape[0] == 1: 86 # fix for stupid wild squeeze from OFA/fairseq. 87 outputs = outputs.unsqueeze(0) 88 if self.config.get("group_image_entities"): 89 logits = outputs.view(batch_size, num_entities, 4) 90 else: 91 # When using OFA without group_image_entities, we always have 1 entity per sample 92 logits = outputs.unsqueeze(1) 93 return self.loss(logits, labels) 94 95 batch_size: int = input_ids.shape[0] 96 inner_input_size: int = input_ids.shape[1] 97 fake_input_batch_size: int = batch_size * inner_input_size 98 99 encoded = self.model.encoder( 100 input_ids.view(fake_input_batch_size, -1), 101 src_lengths=length.view(fake_input_batch_size), 102 patch_images=image.repeat_interleave(inner_input_size, 0), 103 # FIXME build patch_masks the same way patches are built from image. 104 patch_masks=image_mask.repeat_interleave(inner_input_size, 0)[:, 0, 0]) 105 106 inner_output_size: int = decoder_input_ids.shape[1] 107 fake_output_batch_size: int = batch_size * inner_output_size 108 decode_per_encode: int = inner_output_size // inner_input_size 109 decoder_sequence_length: int = decoder_input_ids.shape[2] 110 111 encoded["encoder_out"] = [encoded["encoder_out"][0].repeat_interleave(decode_per_encode, 1)] 112 encoded["encoder_padding_mask"] = [encoded["encoder_padding_mask"][0].repeat_interleave(decode_per_encode, 0)] 113 encoded["position_embeddings"] = [encoded["position_embeddings"][0].repeat_interleave(decode_per_encode, 0)] 114 115 decoded = self.model.decoder( 116 decoder_input_ids.view(fake_output_batch_size, decoder_sequence_length), 117 encoder_out=encoded) 118 lm_logits = self.masked_log_softmax(decoded[0], kwargs.get("decoder_constraint")) 119 120 fake_loss_batch_size: int = fake_output_batch_size * decoder_sequence_length 121 lm_loss = self.loss(lm_logits.view(fake_loss_batch_size, -1), decoder_target.view(fake_loss_batch_size)) 122 lm_loss = lm_loss.view(fake_output_batch_size, decoder_sequence_length) 123 loss = lm_loss 124 125 if self.config.get("label_smoothing"): 126 loss *= (1 - self.config.label_smoothing) 127 mask = kwargs["decoder_constraint"].view(fake_output_batch_size, decoder_sequence_length, -1) 128 num_classes = mask.float().sum(2, keepdim=True) 129 num_classes += (num_classes == 0) 130 uniform = mask.float() / num_classes 131 loss += self.config.label_smoothing * - torch.sum(uniform * lm_logits, 2) 132 133 if self.config.get("encouraging_loss"): 134 bonus = torch.log(-torch.expm1(lm_logits)) 135 log_end = self.config.encouraging_loss 136 if log_end != 1: 137 probs = torch.exp(lm_logits) 138 y_log_end = math.log1p(-log_end) 139 bonus_after_log_end = 1 / (log_end - 1) * (probs - log_end) + y_log_end 140 bonus = torch.where(probs > log_end, bonus_after_log_end, bonus) 141 ec_loss = self.loss(bonus.view(fake_loss_batch_size, -1), decoder_target.view(fake_loss_batch_size)).view(fake_output_batch_size, decoder_sequence_length) 142 143 if self.config.get("label_smoothing"): 144 ec_loss *= (1 - self.config.label_smoothing) 145 mask = kwargs["decoder_constraint"].view(fake_output_batch_size, decoder_sequence_length, -1) 146 num_classes = mask.float().sum(2, keepdim=True) 147 num_classes += (num_classes == 0) 148 uniform = mask.float() / num_classes 149 ec_loss += self.config.label_smoothing * - torch.sum(uniform * bonus, 2) 150 151 loss += ec_loss 152 153 if self.config.get("only_train_answer"): 154 loss *= kwargs["decoder_answer_mask"].view(fake_output_batch_size, decoder_sequence_length) 155 loss = loss.sum(1) 156 if self.config.get("loss_mean_sentence_length"): 157 loss /= decoder_target_attention_mask.view(fake_output_batch_size, decoder_sequence_length).sum(1) 158 loss = loss.mean() 159 160 if self.config.get("only_score_answer"): 161 lm_loss *= kwargs["decoder_answer_mask"].view(fake_output_batch_size, decoder_sequence_length) 162 lm_loss = lm_loss.sum(1) 163 164 if self.config.ofa_task == "vqa_gen" and decode_per_encode == 4: 165 logits = -lm_loss.view(batch_size, decode_per_encode) 166 prediction = torch.nn.functional.softmax(logits, dim=1).unsqueeze(1) 167 elif self.config.ofa_task == "snli_ve" and decode_per_encode in [2, 3]: 168 assert(self.config.get("neither_for_other")) # For another commit 169 assert(self.config.maybe_is not in ["yes", "no"]) # For another commit too =( 170 # logits is of size: batch × label × answer: 171 # - label comes from the encoding of "entity is a hero", "entity is a villain", etc. 172 # - answer comes from the SNLI question answer: "no", "maybe" or "yes" 173 logits = -lm_loss.view(batch_size, inner_input_size, decode_per_encode) 174 175 # FIXME just compare the probabilities of "yes" for now 176 prediction = torch.nn.functional.softmax(logits[:, :, -1], dim=1).unsqueeze(1) 177 else: 178 prediction = None 179 180 return loss, prediction