train.py (15321B)
1 import argparse 2 import contextlib 3 import copy 4 import logging 5 import math 6 import multiprocessing 7 import pathlib 8 import signal 9 import sys 10 import tqdm 11 import types 12 13 # Do this do deal with the badly packaged OFA 14 import mmsrl 15 sys.path.append(mmsrl.__path__[0] + "/../OFA") 16 17 import clip 18 import fairseq.checkpoint_utils 19 import fairseq.dataclass.utils 20 import fairseq.options 21 import fairseq.tasks 22 import fairseq.utils 23 import sklearn.metrics 24 import tasks # OFA 25 import torch 26 import transformers 27 28 from mmsrl.dataset import MemeDataset 29 import mmsrl.batcher 30 import mmsrl.extract_image_features 31 import mmsrl.output 32 import mmsrl.utils 33 34 35 logging.basicConfig(level=logging.INFO) 36 transformers_logger = logging.getLogger("transformers") 37 transformers_logger.setLevel(logging.WARNING) 38 39 40 class Trainer: 41 def __init__(self, config): 42 self.config = config 43 self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") 44 self.prepare_framework() 45 self.prepare_data() 46 self.prepare_model() 47 48 def hook_signals(self) -> None: 49 """ Change the behavior of SIGINT (^C) to change a variable `self.interrupted' before killing the process. """ 50 self.interrupted = False 51 52 def handler(sig: int, frame: types.FrameType) -> None: 53 if multiprocessing.current_process().name != "MainProcess": 54 return 55 56 print("\n\033[31mInterrupted, execution will stop at the next end of validation.\n\033[1mNEXT ^C WILL KILL THE PROCESS!\033[0m\n", file=sys.stderr) 57 self.interrupted = True 58 signal.signal(signal.SIGINT, signal.SIG_DFL) 59 60 signal.signal(signal.SIGINT, handler) 61 62 def get_ofa_config(self): 63 parser = fairseq.options.get_generation_parser() 64 input_args = ["", 65 f"--task={self.config.ofa_task}", 66 "--unnormalized", 67 f"--path={self.config.ofa_path}", 68 "--bpe-dir=OFA/utils/BPE"] 69 args = fairseq.options.parse_args_and_arch(parser, input_args) 70 ofa_config = fairseq.dataclass.utils.convert_namespace_to_omegaconf(args) 71 return ofa_config 72 73 def prepare_feature(self, feature): 74 """ Adapt image and text feature to the model being in float16 or float32 setting. """ 75 if isinstance(feature, torch.Tensor): 76 if self.config.get("half"): 77 feature = feature.to(dtype=torch.float16) 78 feature = feature.to(device=self.device) 79 return feature 80 81 def prepare_framework(self): 82 if self.config.model == "ofa": 83 if self.config.ofa_task == 'vqa_gen': 84 fairseq.tasks.register_task('vqa_gen', tasks.mm_tasks.vqa_gen.VqaGenTask) 85 elif self.config.ofa_task == 'snli_ve': 86 fairseq.tasks.register_task('snli_ve', tasks.mm_tasks.snli_ve.SnliVeTask) 87 else: 88 raise RuntimeError(f"Unknown OFA task {self.config.ofa_task}") 89 ofa_pre_config = self.get_ofa_config() 90 self.ofa_task = fairseq.tasks.setup_task(ofa_pre_config.task) 91 self.ofa_models, self.ofa_config = fairseq.checkpoint_utils.load_model_ensemble( 92 fairseq.utils.split_paths(ofa_pre_config.common_eval.path), 93 task=self.ofa_task 94 ) 95 elif self.config.model in ["clip", "clip_momenta", "generic"]: 96 self.clip_model, self.clip_preprocess = clip.load(self.config.modelname, device=self.device) 97 98 def prepare_data(self): 99 self.epoch = 0 100 self.dataset = {} 101 self.iterator = {} 102 if self.config.model == "ofa": 103 extra_dataset_kwargs = {"ofa_task": self.ofa_task} 104 text_pad: int = self.ofa_task.src_dict.pad() 105 elif self.config.model == "clip": 106 text_pad: int = 0 107 if self.config.features_path: 108 self.image_feature_extractor = mmsrl.extract_image_features.ImageFeatureExtractor(self.config) 109 extra_dataset_kwargs = {"image_processor": self.clip_preprocess, "tokenizer": clip.tokenize, "image_features_loader": self.image_feature_extractor.process_image} 110 else: 111 extra_dataset_kwargs = {"image_processor": self.clip_preprocess, "tokenizer": clip.tokenize} 112 elif self.config.model == "clip_momenta": 113 text_pad: int = 0 114 self.image_feature_extractor = mmsrl.extract_image_features.ImageFeatureExtractor(self.config) 115 extra_dataset_kwargs = {"image_processor": self.clip_preprocess, "tokenizer": clip.tokenize, "image_features_loader": self.image_feature_extractor.process_image} 116 elif self.config.model == "visualbert": 117 self.image_feature_extractor = mmsrl.extract_image_features.ImageFeatureExtractor(self.config) 118 self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.config.modelname) 119 text_pad: int = self.tokenizer.pad_token_id 120 extra_dataset_kwargs = {"image_processor": self.image_feature_extractor.process_image, "tokenizer": self.tokenizer} 121 elif self.config.model == "generic": 122 # tokenizer: 123 extra_dataset_kwargs = {"tokenizer": clip.tokenize, "image_processor": self.clip_preprocess} 124 text_pad: int = 0 125 if self.config.use_visualbert or self.config.features_path: 126 self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.config.visualbert_text_model) 127 self.image_feature_extractor = mmsrl.extract_image_features.ImageFeatureExtractor(self.config) 128 extra_dataset_kwargs.update({"image_features_loader": self.image_feature_extractor.process_image, "vbert_tokenizer": self.tokenizer}) 129 else: 130 self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.config.modelname) 131 text_pad: int = self.tokenizer.pad_token_id 132 extra_dataset_kwargs = {"tokenizer": self.tokenizer} 133 134 self.batcher = mmsrl.batcher.Batcher(self.config, text_pad=text_pad, label_pad=-100, image_pad=0) 135 136 for split in ["train", "val", "test"]: 137 self.dataset[split] = MemeDataset(self.config, split, split!="train", **extra_dataset_kwargs) 138 self.iterator[split] = lambda split=split, trainer=self: tqdm.tqdm( 139 torch.utils.data.DataLoader(dataset=trainer.dataset[split], 140 batch_size=trainer.config["batch_size"], 141 collate_fn=trainer.batcher, 142 num_workers=trainer.config.get("workers", 0), 143 worker_init_fn=trainer.dataset[split].seed, 144 ), desc=f"{split}, epoch: {self.epoch}", leave=False) 145 146 def prepare_model(self) -> None: 147 label_weights = self.dataset['train'].get_label_weights() 148 if self.config.model=="ofa": 149 from mmsrl.ofa import Model 150 model_kwargs = {"ofa_models": self.ofa_models, "ofa_config": self.ofa_config, "ofa_task": self.ofa_task} 151 elif self.config.model=="perceiver": 152 from mmsrl.perceiver import Model 153 model_kwargs = {} 154 elif self.config.model=="clip": 155 from mmsrl.clip import Model 156 model_kwargs = {"model": self.clip_model} 157 elif self.config.model=="clip_momenta": 158 from mmsrl.clip_momenta import Model 159 model_kwargs = {"model": self.clip_model} 160 elif self.config.model == "visualbert": 161 from mmsrl.visualbert import Model 162 model_kwargs = {} 163 elif self.config.model == "generic": 164 from mmsrl.generic_classifier import Model 165 model_kwargs = {"model": self.clip_model} 166 else: 167 from mmsrl.baseline import Model 168 model_kwargs = {} 169 self.model = Model(self.config, label_weights.to(device=self.device), **model_kwargs) 170 if self.config.get("half"): 171 self.model.half() 172 self.model.to(self.device) 173 self.optimizer = torch.optim.AdamW(self.model.parameters(), 174 lr=self.config["learning_rate"], 175 weight_decay=self.config.get("weight_decay", 0)) 176 step_per_epoch = math.ceil(len(self.dataset["train"]) / self.config.batch_size / self.config.accumulation) 177 max_iterations: int = self.config["max_epoch"] * step_per_epoch 178 179 warmup_step = self.config.get("warmup_step", 0) 180 warmup_step = warmup_step * max_iterations if 0 < warmup_step < 1 else warmup_step 181 self.scheduler = transformers.get_scheduler(self.config.get("lr_scheduler", "linear"), self.optimizer, warmup_step, max_iterations) 182 self.scaler = torch.cuda.amp.GradScaler(1) # Used to scale gradients to avoid rounding to 0 when using float16 183 self.amp = torch.cuda.amp.autocast if self.config.amp else contextlib.nullcontext 184 185 def eval(self, split: str, final: bool = False) -> float: 186 with torch.no_grad(), mmsrl.output.Output(self.config, split, self.epoch, final) as output: 187 predictions = [] 188 truth = [] 189 self.model.eval() 190 accuracy_accumulator = 0 191 num_samples = 0 192 loop = self.iterator[split]() 193 for batch in loop: 194 batch = {key: self.prepare_feature(value) for key, value in batch.items()} 195 with self.amp(): # Use mixed-precision float16 196 loss, prediction = self.model(**batch) 197 output(batch, prediction) 198 if prediction is not None and "labels" in batch: 199 predictions.extend(prediction.argmax(2)[batch['entities_mask']].tolist()) 200 truth.extend(batch["labels"][batch['entities_mask']].tolist()) 201 accuracy_accumulator += ((prediction.argmax(2) == batch["labels"])[batch['entities_mask']]).sum().item() 202 num_samples += batch["entities_mask"].sum().item() 203 loop.set_postfix(accuracy=f"{accuracy_accumulator/num_samples:5.3f}") # for the tqdm 204 if predictions and truth: 205 f1 = sklearn.metrics.f1_score(truth, predictions, average="macro") 206 class_precision, class_recall, class_f1, _ = sklearn.metrics.precision_recall_fscore_support(truth, predictions) 207 class_scores = ", ".join(f"{label}[F{class_f1[i]:4.2f}/P{class_precision[i]:4.2f}/R{class_recall[i]:4.2f}]" for i, label in enumerate(mmsrl.utils.LABELS)) 208 print(f"Epoch {self.epoch:3} {split}, Macro-F1: {f1:5.3f}, {class_scores}") 209 return f1 210 return None 211 212 def train_step(self) -> None: 213 """ Apply the gradients to the parameters. """ 214 if self.config.amp: 215 self.scaler.unscale_(self.optimizer) # Gradients are scaled in order to avoid rounding to 0 when casting to float16 216 torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config["max_grad_norm"]) # Clip to gradients to a maximum value (to be robust to outliers) 217 if self.config.amp: 218 self.scaler.step(self.optimizer) # Add the gradient to the parameters 219 else: 220 self.optimizer.step() 221 self.scheduler.step() # Update the learning rate 222 if self.config.amp: 223 self.scaler.update() # Update by how much the gradients should be scaled to avoid rounding to 0 224 self.optimizer.zero_grad() # Set the gradients to 0 to prepare for the next iteration 225 226 def train(self) -> None: 227 best_dev = -math.inf 228 best_dev_epoch = 0 229 best_state_dict = None # Used to save the best model 230 metric_key = None 231 232 self.hook_signals() 233 if not self.config.get("no_initial_valid"): 234 best_dev = self.eval("val", final=True) 235 236 for self.epoch in range(1, 1 + self.config["max_epoch"]): 237 # Accumulate gradient to better estimate it even with small batches 238 if self.interrupted == True: 239 break 240 metric_accumulator = 0 241 num_samples = 0 242 self.model.train() 243 epoch_loop = self.iterator["train"]() 244 for batch_id, batch in enumerate(epoch_loop): 245 if self.interrupted == True: 246 break 247 if batch_id >= self.config.get("batch_per_epoch", math.inf): 248 break 249 batch = {key: self.prepare_feature(value) for key, value in batch.items()} 250 with self.amp(): # Use mixed-precision float16 251 loss, prediction = self.model(**batch) 252 if loss is None: 253 break 254 255 if prediction is not None: 256 metric_accumulator += ((prediction.argmax(2) == batch["labels"])[batch['entities_mask']]).sum().item() 257 metric_key = "accuracy" 258 elif loss is not None: 259 metric_accumulator += loss * batch["entities_mask"].sum().item() 260 metric_key = "loss" 261 num_samples += batch["entities_mask"].sum().item() 262 if metric_key is not None: 263 epoch_loop.set_postfix(**{metric_key: f"{metric_accumulator/num_samples:5.3f}"}) # for the tqdm 264 265 self.scaler.scale(loss).backward() # Accumulate the (scaled) gradients 266 if (1+batch_id) % self.config.get("accumulation", 1) == 0: 267 self.train_step() 268 269 if metric_key is not None: 270 print(f"Epoch {self.epoch:3} train, {metric_key}: {metric_accumulator/num_samples:.4f}") 271 else: 272 print(f"Epoch {self.epoch:3} train, done") 273 if loss is None: 274 break 275 276 if (batch_id+1) % self.config.get("accumulation", 1) != 0: 277 self.train_step() 278 279 dev = self.eval("val") 280 if dev > best_dev: # If dev score improved 281 best_dev = dev 282 best_dev_epoch = self.epoch 283 if self.config.get("save_model"): 284 torch.save(self.model.state_dict(), self.config.save_model) 285 best_state_dict = True 286 else: 287 best_state_dict = copy.deepcopy(self.model.state_dict()) # Save the weights of the model 288 elif self.epoch - best_dev_epoch > self.config.get("patience", 0): # If dev score worsened for several steps (or 1 if patience is 0) 289 break # Early stopping 290 291 self.dataset["train"].compute_subsample(self.epoch) 292 if best_state_dict is not None: # If we improved over random initialization 293 if self.config.get("save_model"): 294 self.model.load_state_dict(torch.load(self.config.save_model)) 295 else: 296 self.model.load_state_dict(best_state_dict) # Load the model with the best dev score 297 for split in ["val", "test"]: 298 self.eval(split, final=True) 299 300 301 if __name__=="__main__": 302 config: mmsrl.utils.dotdict = mmsrl.utils.parse_args() 303 print("\033[33mConfiguration:\033[0m") 304 mmsrl.utils.print_dict(config) 305 Trainer(config).train()