import sys import os import time import logging import torch ## Uncomment the following line to make the code deterministic and use CUBLAS_WORKSPACE_CONFIG=:4096:8 torch.use_deterministic_algorithms(True) import json import numpy as np import random import wandb from omegaconf import OmegaConf, open_dict from os import path from collections import OrderedDict, defaultdict from transformers import get_linear_schedule_with_warmup from transformers import AutoModel, AutoTokenizer from data_utils.utils import load_dataset, load_eval_dataset import pytorch_utils.utils as utils from torch.profiler import profile, record_function, ProfilerActivity from model.entity_ranking_model import EntityRankingModel from model.mention_proposal import MentionProposalModule from data_utils.tensorize_dataset import TensorizeDataset from pytorch_utils.optimization_utils import get_inverse_square_root_decay from utils_evaluate import coref_evaluation from typing import Dict, Union, List, Optional from omegaconf import DictConfig import copy logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO) logger = logging.getLogger() loss_acc_template_dict = { "total": 0.0, "ment_loss": 0.0, "coref": 0.0, "mention_count": 0.0, "processed_docs": 0.0, "ment_correct": 0.0, "ment_total": 0.0, "ment_tp": 0.0, "ment_pp": 0.0, "ment_ap": 0.0, } class Experiment: """Class for training and evaluating coreference models.""" def __init__(self, config: DictConfig): self.config = config print("Seeded: ", config.seed) print("Cuda Available: ", torch.cuda.is_available()) # Whether to train or not self.eval_model: bool = not self.config.train # Initialize dictionary to track key training variables self.train_info = { "val_perf": 0.0, "global_steps": 0, "num_stuck_evals": 0, "peak_memory": 0.0, } self.wandbdata = {} # Initialize model path attributes self.model_path = self.config.paths.model_path self.best_model_path = self.config.paths.best_model_path if not self.eval_model: # Step 1 - Initialize model self._build_model() # Step 2 - Load Data - Data processing choices such as tokenizer will depend on the model self._load_data() # Step 3 - Resume training self._setup_training() # Step 4 - Loading the checkpoint also restores the training metadata self._load_previous_checkpoint() # All set to resume training # But first check if training is remaining if self._is_training_remaining(): self.train() # Perform final evaluation if path.exists(self.best_model_path): # Step 1 - Initialize model self._initialize_best_model() # Step 2 - Load evaluation data self._load_data() # Step 3 - Perform evaluation self.perform_final_eval() else: logger.info("No model accessible!") sys.exit(1) def _build_model(self) -> None: """Constructs the model with given config.""" model_params: DictConfig = self.config.model train_config: DictConfig = self.config.trainer self.model = EntityRankingModel( model_config=model_params, train_config=train_config ) if torch.cuda.is_available(): self.model.cuda(device=self.config.device) # Print model utils.print_model_info(self.model) sys.stdout.flush() def _load_data(self): """Loads and processes the training and evaluation data. Loads the data concerning all the specified datasets for training and eval. The first part of this method loads all the data from the preprocessed jsonline files. In the second half, the loaded data is tensorized for consumption by the model. Apart from loading and processing the data, the method also populates important attributes such as: num_train_docs_map (dict): Dictionary to maintain the number of training docs per dataset which is useful for implementing sampling in joint training. num_training_steps (int): Number of total training steps. eval_per_k_steps (int): Number of gradient updates before each evaluation. """ self.data_iter_map, self.conll_data_dir, self.num_split_docs_map = ( {}, {}, {"train": {}, "dev": {}, "test": {}}, ) raw_data_map = {} max_segment_len: int = self.config.model.doc_encoder.transformer.max_segment_len model_name: str = self.config.model.doc_encoder.transformer.name add_speaker_tokens: bool = self.config.model.doc_encoder.add_speaker_tokens base_data_dir: str = path.abspath(self.config.paths.base_data_dir) # Load data for dataset_name, attributes in self.config.datasets.items(): num_train_docs: Optional[int] = attributes.get("num_train_docs", None) num_dev_docs: Optional[int] = attributes.get("num_dev_docs", None) num_test_docs: Optional[int] = attributes.get("num_test_docs", None) singleton_file: Optional[str] = attributes.get("singleton_file", None) external_md_file: Optional[str] = attributes.get("external_md_file", None) if singleton_file is not None: singleton_file = path.join(base_data_dir, singleton_file) if path.exists(singleton_file): logger.info(f"Singleton file found: {singleton_file}") if external_md_file is not None: external_md_file = path.join(base_data_dir, external_md_file) if path.exists(external_md_file): logger.info( f"External mention detector file found: {external_md_file}" ) # Data directory is a function of dataset name and tokenizer used data_dir = path.join(path.join(base_data_dir, dataset_name), model_name) # Check if speaker tokens are added if add_speaker_tokens: pot_data_dir = path.join( path.join(path.join(base_data_dir, dataset_name)), model_name + "_speaker", ) if path.exists(pot_data_dir): data_dir = pot_data_dir # Datasets such as litbank have cross validation splits if attributes.get("cross_val_split", None) is not None: data_dir = path.join(data_dir, str(attributes.get("cross_val_split"))) logger.info("Data directory: %s" % data_dir) # CoNLL data dir if attributes.get("has_conll", False): conll_dir = path.join( path.join(path.join(base_data_dir, dataset_name)), "conll" ) if attributes.get("cross_val_split", None) is not None: # LitBank like datasets have cross validation splits conll_dir = path.join( conll_dir, str(attributes.get("cross_val_split")) ) if path.exists(conll_dir): self.conll_data_dir[dataset_name] = conll_dir self.num_split_docs_map["train"][dataset_name] = num_train_docs self.num_split_docs_map["dev"][dataset_name] = num_dev_docs self.num_split_docs_map["test"][dataset_name] = num_test_docs if self.eval_model: print("In Eval Model DataLoader") raw_data_map[dataset_name] = load_eval_dataset( data_dir, external_md_file=external_md_file, max_segment_len=max_segment_len, dataset_name=dataset_name, ) else: raw_data_map[dataset_name] = load_dataset( data_dir, singleton_file=singleton_file, num_dev_docs=num_dev_docs, num_test_docs=num_test_docs, max_segment_len=max_segment_len, dataset_name=dataset_name, ) # Tensorize data data_processor = TensorizeDataset( self.model.get_tokenizer(), remove_singletons=(not self.config.keep_singletons), ) if self.eval_model: for split in ["dev", "test"]: self.data_iter_map[split] = {} for dataset in raw_data_map: for split in raw_data_map[dataset]: self.data_iter_map[split][dataset] = data_processor.tensorize_data( raw_data_map[dataset][split], training=False ) else: # Training for split in ["train", "dev", "test"]: self.data_iter_map[split] = {} training = split == "train" for dataset in raw_data_map: self.data_iter_map[split][dataset] = data_processor.tensorize_data( raw_data_map[dataset][split], training=training ) # Estimate number of training steps if self.config.trainer.eval_per_k_steps is None: # Eval steps is 1 epoch (with subsampling) of all the datasets used in joint training self.config.trainer.eval_per_k_steps = sum( self.num_split_docs_map["train"].values() ) self.config.trainer.num_training_steps = ( self.config.trainer.eval_per_k_steps * self.config.trainer.max_evals ) logger.info( f"Number of training steps: {self.config.trainer.num_training_steps}" ) logger.info(f"Eval per k steps: {self.config.trainer.eval_per_k_steps}") def _load_previous_checkpoint(self): """Loads the last checkpoint or best checkpoint.""" # Resume training print("Model Path: ", self.model_path) print("Model Initialised:", torch.cuda.memory_summary(self.config.device)) if path.exists(self.model_path): self.load_model(self.model_path, last_checkpoint=True) logger.info("Model loaded\n") print( "Loaded Model Returned:", torch.cuda.memory_summary(self.config.device) ) else: # Starting training logger.info("Model initialized\n") sys.stdout.flush() def _is_training_remaining(self): """Check if training is done or remaining. There are two cases where we don't resume training: (a) The dev performance has not improved for the allowed patience parameter number of evaluations. (b) Number of gradient updates is already >= Total training steps. Returns: bool: If true, we resume training. Otherwise do final evaluation. """ if self.train_info["num_stuck_evals"] >= self.config.trainer.patience: return False if self.train_info["global_steps"] >= self.config.trainer.num_training_steps: return False return True def _setup_training(self): """Initialize optimizer and bookkeeping variables for training.""" # Dictionary to track key training variables self.train_info = { "val_perf": 0.0, "global_steps": 0, "num_stuck_evals": 0, "peak_memory": 0.0, "max_mem": 0.0, } # Initialize optimizers self._initialize_optimizers() def _initialize_optimizers(self): """Initialize model + optimizer(s). Check if there's a checkpoint in which case we resume from there.""" optimizer_config: DictConfig = self.config.optimizer train_config: DictConfig = self.config.trainer self.optimizer, self.optim_scheduler = {}, {} if torch.cuda.is_available(): # Gradient scaler required for mixed precision training self.scaler = torch.GradScaler("cuda") else: self.scaler = None # Optimizer for clustering params self.optimizer["mem"] = torch.optim.Adam( self.model.get_params()[1], lr=optimizer_config.init_lr, eps=1e-6 ) if optimizer_config.lr_decay == "inv": self.optim_scheduler["mem"] = get_inverse_square_root_decay( self.optimizer["mem"], num_warmup_steps=0 ) else: # No warmup steps for model params self.optim_scheduler["mem"] = get_linear_schedule_with_warmup( self.optimizer["mem"], num_warmup_steps=0, num_training_steps=train_config.num_training_steps, ) if self.config.model.doc_encoder.finetune: # Optimizer for document encoder no_decay = [ "bias", "LayerNorm.weight", ] # No weight decay for bias and layernorm weights encoder_params = self.model.get_params(named=True)[0] grouped_param = [ { "params": [ p for n, p in encoder_params if not any(nd in n for nd in no_decay) ], "lr": optimizer_config.fine_tune_lr, "weight_decay": 1e-2, }, { "params": [ p for n, p in encoder_params if any(nd in n for nd in no_decay) ], "lr": optimizer_config.fine_tune_lr, "weight_decay": 0.0, }, ] self.optimizer["doc"] = torch.optim.AdamW( grouped_param, lr=optimizer_config.fine_tune_lr, eps=1e-6 ) # Scheduler for document encoder num_warmup_steps = int(0.1 * train_config.num_training_steps) if optimizer_config.lr_decay == "inv": self.optim_scheduler["doc"] = get_inverse_square_root_decay( self.optimizer["doc"], num_warmup_steps=num_warmup_steps ) else: self.optim_scheduler["doc"] = get_linear_schedule_with_warmup( self.optimizer["doc"], num_warmup_steps=num_warmup_steps, num_training_steps=train_config.num_training_steps, ) def agg(self, datadepdict): agg_dict = defaultdict(float) for dataset in datadepdict: for key in datadepdict[dataset]: agg_dict[key] += datadepdict[dataset][key] agg_dict["loss_norm"] = ( agg_dict["coref"] / agg_dict["mention_count"] + agg_dict["ment_loss"] / agg_dict["ment_total"] if agg_dict["mention_count"] > 0 else 0 ) agg_dict["ment_acc"] = agg_dict["ment_correct"] / agg_dict["ment_total"] agg_dict["ment_prec"] = ( agg_dict["ment_tp"] / agg_dict["ment_pp"] if agg_dict["ment_pp"] > 0 else 0 ) agg_dict["ment_rec"] = ( agg_dict["ment_tp"] / agg_dict["ment_ap"] if agg_dict["ment_ap"] > 0 else 0 ) agg_dict["ment_f1"] = ( 2 * (agg_dict["ment_prec"] * agg_dict["ment_rec"]) / (agg_dict["ment_prec"] + agg_dict["ment_rec"]) if (agg_dict["ment_prec"] + agg_dict["ment_rec"]) > 0 else 0 ) return agg_dict def train(self) -> None: """Method for training the model. This method implements the training loop. Within the training loop, the model is periodically evaluated on the dev set(s). """ model, optimizer, scheduler, scaler = ( self.model, self.optimizer, self.optim_scheduler, self.scaler, ) model.train() optimizer_config, train_config = self.config.optimizer, self.config.trainer start_time = time.time() eval_time = {"total_time": 0, "num_evals": 0} print("Started Training..") while True: logger.info("Steps done %d" % (self.train_info["global_steps"])) train_data = self.runtime_load_dataset("train") np.random.shuffle(train_data) logger.info("Per epoch training steps: %d" % len(train_data)) logger.info("Per epoch training steps: %d" % len(train_data)) encoder_params, task_params = model.get_params() stat_per_dataset = defaultdict( lambda: copy.deepcopy(loss_acc_template_dict) ) agg_stat = self.agg # Training "epoch" -> May not correspond to actual epoch for cur_document in train_data: def handle_example(document: Dict) -> Union[None, float]: self.train_info["global_steps"] += 1 for key in optimizer: optimizer[key].zero_grad() loss_dict: Dict = model.forward_training(document) total_loss = loss_dict["total"] if total_loss is None or torch.isnan(total_loss): print("Problem with Loss. Should not occur often") return None total_loss.backward() # Gradient clipping try: for name_ind, param_group in enumerate( [encoder_params, task_params] ): torch.nn.utils.clip_grad_norm_( param_group, optimizer_config.max_gradient_norm, error_if_nonfinite=True, ) except RuntimeError: print("Non Finite Gradient") return None for key in optimizer: self.wandbdata[key + "_lr"] = scheduler[key].get_last_lr()[0] for key in optimizer: optimizer[key].step() scheduler[key].step() loss_dict_items = {} for key in loss_dict: loss_dict_items[key] = loss_dict[key].item() dataset_name = document["dataset_name"] # print(f"Total loss {cur_document['doc_key']}: {total_loss.item()}") for key in loss_dict_items: stat_per_dataset[dataset_name][key] += loss_dict_items[key] stat_per_dataset[dataset_name]["processed_docs"] += 1 return total_loss.item() loss = handle_example(cur_document) if self.train_info["global_steps"] % train_config.log_frequency == 0: max_mem = ( ( torch.cuda.max_memory_allocated(self.config.device) / (1024**3) ) if torch.cuda.is_available() else 0.0 ) if self.train_info.get("max_mem", 0.0) < max_mem: self.train_info["max_mem"] = max_mem if loss is not None: logger.info( "{} {:.3f} Max mem {:.1f} GB".format( cur_document["doc_key"], loss, max_mem, ) ) sys.stdout.flush() if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() if train_config.eval_per_k_steps and ( self.train_info["global_steps"] % train_config.eval_per_k_steps == 0 ): print("Eval needs to be done here") coref_dict = {} print(stat_per_dataset) if self.config.use_wandb: self._wandb_log( split="train", stat_per_dataset=stat_per_dataset, agg_stat=agg_stat, coref_dict=coref_dict, step=self.train_info["global_steps"] // train_config.eval_per_k_steps, ) stat_per_dataset = defaultdict( lambda: copy.deepcopy(loss_acc_template_dict) ) macro_fscore = self.periodic_model_eval() model.train() # Get elapsed time elapsed_time = time.time() - start_time start_time = time.time() logger.info( "Steps: %d, Micro F1: %.1f, Max Micro F1: %.1f, Time: %.2f" % ( self.train_info["global_steps"], macro_fscore, self.train_info["val_perf"], elapsed_time, ) ) # Check stopping criteria if not self._is_training_remaining(): break # Check stopping criteria if not self._is_training_remaining(): break logger.handlers[0].flush() def runtime_load_dataset(self, split): # Shuffle and load the training data data = [] for dataset, dataset_data in self.data_iter_map[split].items(): np.random.shuffle( dataset_data ) ### Commenting this so that we can have a deterministic training if self.num_split_docs_map[split].get(dataset, None) is not None: # Subsampling the data - This is useful in joint training logger.info( f"{dataset}: Subsampled {self.num_split_docs_map[split].get(dataset)}" ) random_indices = range(self.num_split_docs_map[split].get(dataset)) data += [dataset_data[idx] for idx in random_indices] else: data += dataset_data return data def _wandb_log(self, split, stat_per_dataset, agg_stat, coref_dict, step=None): for dataset_name in stat_per_dataset: for metric_vals in stat_per_dataset[dataset_name]: wandb.log( data={ f"{split}/{dataset_name}/{metric_vals}": stat_per_dataset[ dataset_name ][metric_vals] }, step=step, ) if stat_per_dataset[dataset_name]["mention_count"] > 0.0: ment_prec = ( stat_per_dataset[dataset_name]["ment_tp"] / stat_per_dataset[dataset_name]["ment_pp"] if stat_per_dataset[dataset_name]["ment_pp"] > 0 else 0 ) ment_rec = ( stat_per_dataset[dataset_name]["ment_tp"] / stat_per_dataset[dataset_name]["ment_ap"] if stat_per_dataset[dataset_name]["ment_ap"] > 0 else 0 ) ment_f1 = ( 2 * (ment_prec * ment_rec) / (ment_prec + ment_rec) if (ment_prec + ment_rec) > 0 else 0 ) wandb.log( data={ f"{split}/{dataset_name}/loss_norm": stat_per_dataset[ dataset_name ]["coref"] / stat_per_dataset[dataset_name]["mention_count"] + stat_per_dataset[dataset_name]["ment_loss"] / stat_per_dataset[dataset_name]["ment_total"], f"{split}/{dataset_name}/ment_acc": stat_per_dataset[ dataset_name ]["ment_correct"] / stat_per_dataset[dataset_name]["ment_total"], f"{split}/{dataset_name}/ment_prec": ment_prec, f"{split}/{dataset_name}/ment_rec": ment_rec, f"{split}/{dataset_name}/ment_f1": ment_f1, }, step=step, ) else: print("No mentions processed. Should not occur many times.") if agg_stat: for metric in agg_stat(stat_per_dataset): wandb.log( data={f"{split}/{metric}": agg_stat(stat_per_dataset)[metric]}, step=step, ) for dataset in coref_dict: for key in coref_dict[dataset]: # Log result for individual metrics if isinstance(coref_dict[dataset][key], dict): wandb.log( data={ f"{split}/{dataset}/{key}": coref_dict[dataset][key].get( "fscore", 0.0 ) }, step=step, ) # Log the overall F-score wandb.log( data={ f"{split}/{dataset}/CoNLL": coref_dict[dataset].get("fscore", 0.0) }, step=step, ) wandb.log( data={ f"{split}/{dataset}/Micro-F1": coref_dict[dataset].get( "f1_micro", 0.0 ) }, step=step, ) wandb.log( data={ f"{split}/{dataset}/Macro-F1": coref_dict[dataset].get( "f1_macro", 0.0 ) }, step=step, ) wandb.log(data=self.wandbdata, step=step) @torch.no_grad() def periodic_model_eval(self) -> float: """Method for evaluating and saving the model during the training loop. Returns: float: Average CoNLL F-score over all the development sets of datasets. """ self.model.eval() ## Dev Loss Calculations: dev_data = self.runtime_load_dataset("dev") np.random.shuffle(dev_data) stat_per_dataset = defaultdict(lambda: copy.deepcopy(loss_acc_template_dict)) agg_stat = self.agg for cur_document in dev_data: def handle_example(document: Dict) -> Union[None, float]: loss_dict: Dict = self.model.forward_training(document) total_loss = loss_dict["total"] if total_loss is None or torch.isnan(total_loss): print("Problem with Loss. Should not occur many times") return None loss_dict_items = {} for key in loss_dict: loss_dict_items[key] = loss_dict[key].item() dataset_name = document["dataset_name"] for key in loss_dict_items: stat_per_dataset[dataset_name][key] += loss_dict_items[key] stat_per_dataset[dataset_name]["processed_docs"] += 1 return total_loss.item() loss = handle_example(cur_document) if loss is None: continue # Dev performance coref_dict = {} train_config = self.config.trainer for dataset in self.data_iter_map["dev"]: for go in [False]: for tf in [False]: result_dict = coref_evaluation( self.config, self.model, self.data_iter_map, dataset, teacher_force=tf, gold_mentions=go, _iter="_" + str( self.train_info["global_steps"] // train_config.eval_per_k_steps ), conll_data_dir=self.conll_data_dir, ) coref_dict[dataset] = result_dict if self.config.use_wandb: self._wandb_log( split="dev", stat_per_dataset=stat_per_dataset, agg_stat=agg_stat, coref_dict=coref_dict, step=self.train_info["global_steps"] // train_config.eval_per_k_steps, ) # Calculate Mean F-score fscore = sum([coref_dict[dataset]["fscore"] for dataset in coref_dict]) / len( coref_dict ) micro_fscore = sum( [coref_dict[dataset]["f1_micro"] for dataset in coref_dict] ) / len(coref_dict) macro_fscore = sum( [coref_dict[dataset]["f1_macro"] for dataset in coref_dict] ) / len(coref_dict) logger.info( "Avg Macro F1: %.1f, Max Micro F1: %.1f" % (macro_fscore, self.train_info["val_perf"]) ) logger.info("Avg Macro F1: %.1f" % (macro_fscore)) # Update model if dev performance improves if macro_fscore > self.train_info["val_perf"]: # Update training bookkeeping variables self.train_info["num_stuck_evals"] = 0 self.train_info["val_perf"] = macro_fscore # Save the best model logger.info("Saving best model") self.save_model(self.best_model_path, last_checkpoint=False) else: self.train_info["num_stuck_evals"] += 1 # Save model if self.config.trainer.to_save_model: self.save_model(self.model_path, last_checkpoint=True) # Go back to training mode self.model.train() return macro_fscore @torch.no_grad() def perform_final_eval(self) -> None: """Method to evaluate the model after training has finished.""" self.model.eval() base_output_dict = OmegaConf.to_container(self.config) perf_summary = {"best_perf": self.train_info["val_perf"]} if self.config.paths.model_dir: perf_summary["model_dir"] = path.normpath(self.config.paths.model_dir) logger.info( "Max training memory: %.1f GB" % self.train_info.get("max_mem", 0.0) ) logger.info("Validation performance: %.1f" % self.train_info["val_perf"]) perf_file_dict = {} dataset_output_dict = {} for split in ["dev", "test"]: perf_summary[split] = {} logger.info("\n") logger.info("%s" % split.capitalize()) coref_dict = {} for dataset in self.data_iter_map.get(split, []): dataset_dir = path.join(self.config.paths.model_dir, dataset) if not path.exists(dataset_dir): os.makedirs(dataset_dir) if dataset not in dataset_output_dict: dataset_output_dict[dataset] = {} if dataset not in perf_file_dict: perf_file_dict[dataset] = path.join(dataset_dir, f"perf.json") print("Dataset Name:", self.config.datasets[dataset].name) logger.info("Dataset: %s\n" % self.config.datasets[dataset].name) for go in [False]: for tf in [False]: result_dict = coref_evaluation( self.config, self.model, self.data_iter_map, dataset=dataset, split=split, teacher_force=tf, gold_mentions=go, final_eval=True, conll_data_dir=self.conll_data_dir, ) coref_dict[dataset] = result_dict dataset_output_dict[dataset][split] = result_dict perf_summary[split][dataset] = result_dict["f1_micro"] if self.config.use_wandb: self._wandb_log( split=split, stat_per_dataset={}, agg_stat=None, coref_dict=coref_dict, step=None, ) sys.stdout.flush() for dataset, output_dict in dataset_output_dict.items(): perf_file = perf_file_dict[dataset] json.dump(output_dict, open(perf_file, "w"), indent=2) logger.info("Final performance summary at %s" % path.abspath(perf_file)) summary_file = path.join(self.config.paths.model_dir, "perf.json") json.dump(perf_summary, open(summary_file, "w"), indent=2) logger.info("Performance summary file: %s" % path.abspath(summary_file)) def _initialize_best_model(self): checkpoint = torch.load( self.best_model_path, map_location="cpu", ) config = checkpoint["config"] ## Due to version changes -- these changes are necessary # if if self.config.get("override_encoder", False): model_config = config.model print(type(self.config.model.doc_encoder.transformer)) print(self.config.model.doc_encoder.transformer) model_config.doc_encoder.transformer = ( self.config.model.doc_encoder.transformer ) # Override memory # For e.g., can test with a different bounded memory size if self.config.get("override_memory", False): model_config = config.model model_config.memory = self.config.model.memory with open_dict(config): print("Config change") config.model.mention_params.ext_ment = ( self.config.model.mention_params.ext_ment ) config = utils.fill_missing_configs(config, self.config) print("Type: ", config.model.memory.type) self.config.model = config.model self.train_info = checkpoint["train_info"] if self.config.model.doc_encoder.finetune: # Load the document encoder params if encoder is finetuned doc_encoder_dir = path.join( path.dirname(self.best_model_path), self.config.paths.doc_encoder_dirname, ) if path.exists(doc_encoder_dir): logger.info( "Loading document encoder from %s" % path.abspath(doc_encoder_dir) ) config.model.doc_encoder.transformer.model_str = doc_encoder_dir self.model = EntityRankingModel(config.model, config.trainer) # Document encoder parameters will be loaded via the huggingface initialization self.model.load_state_dict(checkpoint["model"], strict=False) if torch.cuda.is_available(): self.model.cuda(device=self.config.device) def load_model(self, location: str, last_checkpoint=True) -> None: """Load model from given location. Args: location: str Location of checkpoint last_checkpoint: bool Whether the checkpoint is the last one saved or not. If false, don't load optimizers, schedulers, and other training variables. """ checkpoint = torch.load(location, map_location="cpu") logger.info("Loading model from %s" % path.abspath(location)) # self.config = checkpoint["config"] ## Commented out so that it does not load the config of the trained model. Removed comment self.model.load_state_dict( checkpoint["model"], strict=False ) ## No encoder in this model so strict=False is compulsary. No other weight missing. Checked # self.train_info = checkpoint["train_info"] ## No train info transfer too. ## Transferring if self.config.model.doc_encoder.finetune: # Load the document encoder params if encoder is finetuned doc_encoder_dir = path.join( path.dirname(location), self.config.paths.doc_encoder_dirname ) logger.info( "Loading document encoder from %s" % path.abspath(doc_encoder_dir) ) # Load the encoder self.model.mention_proposer.doc_encoder.lm_encoder = ( AutoModel.from_pretrained(pretrained_model_name_or_path=doc_encoder_dir) ) self.model.mention_proposer.doc_encoder.tokenizer = ( AutoTokenizer.from_pretrained( pretrained_model_name_or_path=doc_encoder_dir, clean_up_tokenization_spaces=True, ) ) if self.model.mention_proposer.doc_encoder.config.finetune: self.model.mention_proposer.doc_encoder.lm_encoder.gradient_checkpointing_enable() if torch.cuda.is_available(): self.model.cuda(device=self.config.device) print("Loaded Model:", torch.cuda.memory_summary()) print( "Gradient checkpointing enabled? ", torch.autograd.grad_checkpoint_enabled() ) del checkpoint torch.cuda.empty_cache() def save_model(self, location: os.PathLike, last_checkpoint=True) -> None: """Save model. Args: location: Location of checkpoint last_checkpoint: Whether the checkpoint is the last one saved or not. If false, don't save optimizers and schedulers which take up a lot of space. """ model_state_dict = OrderedDict(self.model.state_dict()) doc_encoder_state_dict = {} # Separate the doc_encoder state dict # We will save the model in two parts: # (a) Doc encoder parameters - Useful for final upload to huggingface # (b) Rest of the model parameters, optimizers, schedulers, and other bookkeeping variables for key in self.model.state_dict(): if "lm_encoder." in key: doc_encoder_state_dict[key] = model_state_dict[key] del model_state_dict[key] # Save the document encoder params if self.config.model.doc_encoder.finetune: doc_encoder_dir = path.join( path.dirname(location), self.config.paths.doc_encoder_dirname ) if not path.exists(doc_encoder_dir): os.makedirs(doc_encoder_dir) logger.info(f"Encoder saved at {path.abspath(doc_encoder_dir)}") # Save the encoder self.model.mention_proposer.doc_encoder.lm_encoder.save_pretrained( save_directory=doc_encoder_dir, save_config=True ) # Save the tokenizer self.model.mention_proposer.doc_encoder.tokenizer.save_pretrained( doc_encoder_dir ) save_dict = { "train_info": self.train_info, "model": model_state_dict, "rng_state": torch.get_rng_state(), "np_rng_state": np.random.get_state(), "config": self.config, } if self.scaler is not None: save_dict["scaler"] = self.scaler.state_dict() if last_checkpoint: # For last checkpoint save the optimizer and scheduler states as well save_dict["optimizer"] = {} save_dict["scheduler"] = {} param_groups: List[str] = ( ["mem", "doc"] if self.config.model.doc_encoder.finetune else ["mem"] ) for param_group in param_groups: save_dict["optimizer"][param_group] = self.optimizer[ param_group ].state_dict() save_dict["scheduler"][param_group] = self.optim_scheduler[ param_group ].state_dict() torch.save(save_dict, location) logger.info(f"Model saved at: {path.abspath(location)}")