def _evaluate( self, model: EmmentalModel, dataloaders: List[EmmentalDataLoader], split: Union[List[str], str], ) -> Dict[str, float]: r"""Evaluate the model. Args: model(EmmentalModel): The model to evaluate. dataloaders(List[EmmentalDataLoader]): The data to evaluate. split(str): The split to evaluate. Returns: dict: The score dict. """ if not isinstance(split, list): valid_split = [split] else: valid_split = split valid_dataloaders = [ dataloader for dataloader in dataloaders if dataloader.split in valid_split ] return model.score(valid_dataloaders)
def main(args): # Initialize Emmental config = parse_args_to_config(args) emmental.init(log_dir=config["meta_config"]["log_path"], config=config) # Log configuration into files cmd_msg = " ".join(sys.argv) logger.info(f"COMMAND: {cmd_msg}") write_to_file(f"{emmental.Meta.log_path}/cmd.txt", cmd_msg) logger.info(f"Config: {emmental.Meta.config}") write_to_file(f"{emmental.Meta.log_path}/config.txt", emmental.Meta.config) # Create dataloaders dataloaders = get_dataloaders(args) # Assign transforms to dataloaders aug_dataloaders = [] if args.augment_policy: for idx in range(len(dataloaders)): if dataloaders[idx].split in args.train_split: dataloaders[idx].dataset.transform_cls = Augmentation( args=args) config["learner_config"]["task_scheduler_config"][ "task_scheduler"] = AugScheduler(augment_k=args.augment_k, enlarge=args.augment_enlarge) emmental.Meta.config["learner_config"]["task_scheduler_config"][ "task_scheduler"] = config["learner_config"]["task_scheduler_config"][ "task_scheduler"] # Create tasks model = EmmentalModel(name=f"{args.task}_task") model.add_task(create_task(args)) # Set cudnn benchmark cudnn.benchmark = True # Load the best model from the pretrained model if config["model_config"]["model_path"] is not None: model.load(config["model_config"]["model_path"]) if args.train: emmental_learner = EmmentalLearner() emmental_learner.learn(model, dataloaders + aug_dataloaders) # Remove all extra augmentation policy for idx in range(len(dataloaders)): dataloaders[idx].dataset.transform_cls = None scores = model.score(dataloaders) # Save metrics and models logger.info(f"Metrics: {scores}") scores["log_path"] = emmental.Meta.log_path write_to_json_file(f"{emmental.Meta.log_path}/metrics.txt", scores) model.save(f"{emmental.Meta.log_path}/last_model.pth")
def main(args): # Ensure that global state is fresh Meta.reset() # Initialize Emmental config = parse_arg_to_config(args) emmental.init(config["meta_config"]["log_path"], config=config) # Save command line argument into file cmd_msg = " ".join(sys.argv) logger.info(f"COMMAND: {cmd_msg}") write_to_file(Meta.log_path, "cmd.txt", cmd_msg) # Save Emmental config into file logger.info(f"Config: {Meta.config}") write_to_file(Meta.log_path, "config.txt", Meta.config) Meta.config["learner_config"]["global_evaluation_metric_dict"] = { f"model/SuperGLUE/{split}/score": partial(superglue_scorer, split=split) for split in ["val"] } # Construct dataloaders and tasks and load slices dataloaders = [] tasks = [] for task_name in args.task: task_dataloaders = get_dataloaders( data_dir=args.data_dir, task_name=task_name, splits=["train", "val", "test"], max_sequence_length=args.max_sequence_length, max_data_samples=args.max_data_samples, tokenizer_name=args.bert_model, batch_size=args.batch_size, augment=args.augmentations, ) task = models.model[task_name]( args.bert_model, last_hidden_dropout_prob=args.last_hidden_dropout_prob) if args.slices: logger.info("Initializing task-specific slices") slice_func_dict = slicing.slice_func_dict[task_name] # Include general purpose slices if args.general_slices: logger.info("Including general slices") slice_func_dict.update(slicing.slice_func_dict["general"]) task_dataloaders = slicing.add_slice_labels( task_name, task_dataloaders, slice_func_dict) slice_tasks = slicing.add_slice_tasks(task_name, task, slice_func_dict, args.slice_hidden_dim) tasks.extend(slice_tasks) else: tasks.append(task) dataloaders.extend(task_dataloaders) # Build Emmental model model = EmmentalModel(name="SuperGLUE", tasks=tasks) # Load pretrained model if necessary if Meta.config["model_config"]["model_path"]: model.load(Meta.config["model_config"]["model_path"]) # Training if args.train: emmental_learner = EmmentalLearner() emmental_learner.learn(model, dataloaders) # If model is slice-aware, slice scores will be calculated from slice heads # If model is not slice-aware, manually calculate performance on slices if not args.slices: slice_func_dict = {} slice_keys = args.task if args.general_slices: slice_keys.append("general") for k in slice_keys: slice_func_dict.update(slicing.slice_func_dict[k]) scores = slicing.score_slices(model, dataloaders, args.task, slice_func_dict) else: scores = model.score(dataloaders) # Save metrics into file logger.info(f"Metrics: {scores}") write_to_file(Meta.log_path, "metrics.txt", scores) # Save best metrics into file if args.train: logger.info( f"Best metrics: " f"{emmental_learner.logging_manager.checkpointer.best_metric_dict}" ) write_to_file( Meta.log_path, "best_metrics.txt", emmental_learner.logging_manager.checkpointer.best_metric_dict, ) # Save submission file for task_name in args.task: dataloaders = [d for d in dataloaders if d.split == "test"] assert len(dataloaders) == 1 filepath = os.path.join(Meta.log_path, f"{task_name}.jsonl") make_submission_file(model, dataloaders[0], task_name, filepath)
def main(args): # Initialize Emmental config = parse_args_to_config(args) emmental.init(log_dir=config["meta_config"]["log_path"], config=config) # Log configuration into files cmd_msg = " ".join(sys.argv) logger.info(f"COMMAND: {cmd_msg}") write_to_file(f"{emmental.Meta.log_path}/cmd.txt", cmd_msg) logger.info(f"Config: {emmental.Meta.config}") write_to_file(f"{emmental.Meta.log_path}/config.txt", emmental.Meta.config) # Create dataloaders dataloaders = get_dataloaders(args) config["learner_config"]["task_scheduler_config"][ "task_scheduler"] = AugScheduler(augment_k=args.augment_k, enlarge=args.augment_enlarge) emmental.Meta.config["learner_config"]["task_scheduler_config"][ "task_scheduler"] = config["learner_config"]["task_scheduler_config"][ "task_scheduler"] # Specify parameter group for Adam BERT def grouped_parameters(model): no_decay = ["bias", "LayerNorm.weight"] return [ { "params": [ p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) ], "weight_decay": emmental.Meta.config["learner_config"]["optimizer_config"] ["l2"], }, { "params": [ p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) ], "weight_decay": 0.0, }, ] emmental.Meta.config["learner_config"]["optimizer_config"][ "parameters"] = grouped_parameters # Create tasks model = EmmentalModel(name=f"{args.task}_task") model.add_task(create_task(args)) # Load the best model from the pretrained model if config["model_config"]["model_path"] is not None: model.load(config["model_config"]["model_path"]) if args.train: emmental_learner = EmmentalLearner() emmental_learner.learn(model, dataloaders) # Remove all extra augmentation policy for idx in range(len(dataloaders)): dataloaders[idx].dataset.transform_cls = None dataloaders[idx].dataset.k = 1 scores = model.score(dataloaders) # Save metrics and models logger.info(f"Metrics: {scores}") scores["log_path"] = emmental.Meta.log_path write_to_json_file(f"{emmental.Meta.log_path}/metrics.txt", scores) model.save(f"{emmental.Meta.log_path}/last_model.pth")
task_name: create_task( task_name, args, datasets[task_name]["nclasses"], emb_layer ) for task_name in args.task } model = EmmentalModel(name="TC_task") if Meta.config["model_config"]["model_path"]: model.load(Meta.config["model_config"]["model_path"]) else: for task_name, task in tasks.items(): model.add_task(task) emmental_learner = EmmentalLearner() emmental_learner.learn(model, dataloaders) scores = model.score(dataloaders) logger.info(f"Metrics: {scores}") write_to_json_file(f"{Meta.log_path}/metrics.txt", scores) if args.checkpointing: logger.info( f"Best metrics: " f"{emmental_learner.logging_manager.checkpointer.best_metric_dict}" ) write_to_file( f"{Meta.log_path}/best_metrics.txt", emmental_learner.logging_manager.checkpointer.best_metric_dict, )
def test_e2e(caplog): """Run an end-to-end test.""" caplog.set_level(logging.INFO) dirpath = "temp_test_e2e" use_exact_log_path = False Meta.reset() emmental.init(dirpath, use_exact_log_path=use_exact_log_path) config = { "meta_config": { "seed": 0 }, "learner_config": { "n_epochs": 3, "optimizer_config": { "lr": 0.01, "grad_clip": 100 }, }, "logging_config": { "counter_unit": "epoch", "evaluation_freq": 1, "writer_config": { "writer": "tensorboard", "verbose": True }, "checkpointing": True, "checkpointer_config": { "checkpoint_path": None, "checkpoint_freq": 1, "checkpoint_metric": { "model/all/train/loss": "min" }, "checkpoint_task_metrics": None, "checkpoint_runway": 1, "checkpoint_all": False, "clear_intermediate_checkpoints": True, "clear_all_checkpoints": True, }, }, } emmental.Meta.update_config(config) # Generate synthetic data N = 500 X = np.random.random((N, 2)) * 2 - 1 Y1 = (X[:, 0] > X[:, 1] + 0.25).astype(int) Y2 = (X[:, 0] > X[:, 1] + 0.2).astype(int) X = [torch.Tensor(X[i]) for i in range(N)] # Create dataset and dataloader X_train, X_dev, X_test = ( X[:int(0.8 * N)], X[int(0.8 * N):int(0.9 * N)], X[int(0.9 * N):], ) Y1_train, Y1_dev, Y1_test = ( torch.tensor(Y1[:int(0.8 * N)]), torch.tensor(Y1[int(0.8 * N):int(0.9 * N)]), torch.tensor(Y1[int(0.9 * N):]), ) Y2_train, Y2_dev, Y2_test = ( torch.tensor(Y2[:int(0.8 * N)]), torch.tensor(Y2[int(0.8 * N):int(0.9 * N)]), torch.tensor(Y2[int(0.9 * N):]), ) train_dataset1 = EmmentalDataset(name="synthetic", X_dict={"data": X_train}, Y_dict={"label1": Y1_train}) train_dataset2 = EmmentalDataset(name="synthetic", X_dict={"data": X_train}, Y_dict={"label2": Y2_train}) dev_dataset1 = EmmentalDataset(name="synthetic", X_dict={"data": X_dev}, Y_dict={"label1": Y1_dev}) dev_dataset2 = EmmentalDataset(name="synthetic", X_dict={"data": X_dev}, Y_dict={"label2": Y2_dev}) test_dataset1 = EmmentalDataset(name="synthetic", X_dict={"data": X_test}, Y_dict={"label1": Y1_test}) test_dataset2 = EmmentalDataset(name="synthetic", X_dict={"data": X_test}, Y_dict={"label2": Y2_test}) task_to_label_dict = {"task1": "label1"} train_dataloader1 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=train_dataset1, split="train", batch_size=10, ) dev_dataloader1 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=dev_dataset1, split="valid", batch_size=10, ) test_dataloader1 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=test_dataset1, split="test", batch_size=10, ) task_to_label_dict = {"task2": "label2"} train_dataloader2 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=train_dataset2, split="train", batch_size=10, ) dev_dataloader2 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=dev_dataset2, split="valid", batch_size=10, ) test_dataloader2 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=test_dataset2, split="test", batch_size=10, ) # Create task def ce_loss(task_name, immediate_ouput_dict, Y, active): module_name = f"{task_name}_pred_head" return F.cross_entropy(immediate_ouput_dict[module_name][0][active], (Y.view(-1))[active]) def output(task_name, immediate_ouput_dict): module_name = f"{task_name}_pred_head" return F.softmax(immediate_ouput_dict[module_name][0], dim=1) task_metrics = {"task1": ["accuracy"], "task2": ["accuracy", "roc_auc"]} tasks = [ EmmentalTask( name=task_name, module_pool=nn.ModuleDict({ "input_module": nn.Linear(2, 8), f"{task_name}_pred_head": nn.Linear(8, 2), }), task_flow=[ { "name": "input", "module": "input_module", "inputs": [("_input_", "data")], }, { "name": f"{task_name}_pred_head", "module": f"{task_name}_pred_head", "inputs": [("input", 0)], }, ], loss_func=partial(ce_loss, task_name), output_func=partial(output, task_name), scorer=Scorer(metrics=task_metrics[task_name]), ) for task_name in ["task1", "task2"] ] # Build model mtl_model = EmmentalModel(name="all", tasks=tasks) # Create learner emmental_learner = EmmentalLearner() # Learning emmental_learner.learn( mtl_model, [ train_dataloader1, train_dataloader2, dev_dataloader1, dev_dataloader2 ], ) test1_score = mtl_model.score(test_dataloader1) test2_score = mtl_model.score(test_dataloader2) assert test1_score["task1/synthetic/test/accuracy"] >= 0.7 assert (test1_score["model/all/test/macro_average"] == test1_score["task1/synthetic/test/accuracy"]) assert test2_score["task2/synthetic/test/accuracy"] >= 0.7 assert test2_score["task2/synthetic/test/roc_auc"] >= 0.7 shutil.rmtree(dirpath)
def main( conn_string, max_docs=float("inf"), parse=False, first_time=False, gpu=None, parallel=4, log_dir=None, verbose=False, ): if not log_dir: log_dir = "logs" if verbose: level = logging.INFO else: level = logging.WARNING dirname = os.path.dirname(os.path.abspath(__file__)) init_logging(log_dir=os.path.join(dirname, log_dir), level=level) session = Meta.init(conn_string).Session() # Parsing logger.info(f"Starting parsing...") start = timer() docs, train_docs, dev_docs, test_docs = parse_dataset( session, dirname, first_time=first_time, parallel=parallel, max_docs=max_docs ) end = timer() logger.warning(f"Parse Time (min): {((end - start) / 60.0):.1f}") logger.info(f"# of train Documents: {len(train_docs)}") logger.info(f"# of dev Documents: {len(dev_docs)}") logger.info(f"# of test Documents: {len(test_docs)}") logger.info(f"Documents: {session.query(Document).count()}") logger.info(f"Sections: {session.query(Section).count()}") logger.info(f"Paragraphs: {session.query(Paragraph).count()}") logger.info(f"Sentences: {session.query(Sentence).count()}") logger.info(f"Figures: {session.query(Figure).count()}") start = timer() Thumbnails = mention_subclass("Thumbnails") thumbnails_img = MentionFigures() class HasFigures(_Matcher): def _f(self, m): file_path = "" for prefix in [ f"{dirname}/data/train/html/", f"{dirname}/data/dev/html/", f"{dirname}/data/test/html/", ]: if os.path.exists(prefix + m.figure.url): file_path = prefix + m.figure.url if file_path == "": return False img = Image.open(file_path) width, height = img.size min_value = min(width, height) return min_value > 50 mention_extractor = MentionExtractor( session, [Thumbnails], [thumbnails_img], [HasFigures()], parallelism=parallel ) if first_time: mention_extractor.apply(docs) logger.info("Total Mentions: {}".format(session.query(Mention).count())) ThumbnailLabel = candidate_subclass("ThumbnailLabel", [Thumbnails]) candidate_extractor = CandidateExtractor( session, [ThumbnailLabel], throttlers=[None], parallelism=parallel ) if first_time: candidate_extractor.apply(train_docs, split=0) candidate_extractor.apply(dev_docs, split=1) candidate_extractor.apply(test_docs, split=2) train_cands = candidate_extractor.get_candidates(split=0) # Sort the dev_cands, which are used for training, for deterministic behavior dev_cands = candidate_extractor.get_candidates(split=1, sort=True) test_cands = candidate_extractor.get_candidates(split=2) end = timer() logger.warning(f"Candidate Extraction Time (min): {((end - start) / 60.0):.1f}") logger.info("Total train candidate:\t{}".format(len(train_cands[0]))) logger.info("Total dev candidate:\t{}".format(len(dev_cands[0]))) logger.info("Total test candidate:\t{}".format(len(test_cands[0]))) fin = open(f"{dirname}/data/ground_truth.txt", "r") gt = set() for line in fin: gt.add("::".join(line.lower().split())) fin.close() # Labeling start = timer() def LF_gt_label(c): doc_file_id = ( f"{c[0].context.figure.document.name.lower()}.pdf::" f"{os.path.basename(c[0].context.figure.url.lower())}" ) return TRUE if doc_file_id in gt else FALSE gt_dev = [LF_gt_label(cand) for cand in dev_cands[0]] gt_test = [LF_gt_label(cand) for cand in test_cands[0]] end = timer() logger.warning(f"Supervision Time (min): {((end - start) / 60.0):.1f}") batch_size = 64 input_size = 224 K = 2 emmental.init(log_dir=Meta.log_path, config=emmental_config) emmental.Meta.config["learner_config"]["task_scheduler_config"][ "task_scheduler" ] = DauphinScheduler(augment_k=K, enlarge=1) train_dataset = ThumbnailDataset( "Thumbnail", dev_cands[0], gt_dev, "train", prob_label=True, prefix=f"{dirname}/data/dev/html/", input_size=input_size, transform_cls=Augmentation(2), k=K, ) val_dataset = ThumbnailDataset( "Thumbnail", dev_cands[0], gt_dev, "valid", prob_label=False, prefix=f"{dirname}/data/dev/html/", input_size=input_size, k=1, ) test_dataset = ThumbnailDataset( "Thumbnail", test_cands[0], gt_test, "test", prob_label=False, prefix=f"{dirname}/data/test/html/", input_size=input_size, k=1, ) dataloaders = [] dataloaders.append( EmmentalDataLoader( task_to_label_dict={"Thumbnail": "labels"}, dataset=train_dataset, split="train", shuffle=True, batch_size=batch_size, num_workers=1, ) ) dataloaders.append( EmmentalDataLoader( task_to_label_dict={"Thumbnail": "labels"}, dataset=val_dataset, split="valid", shuffle=False, batch_size=batch_size, num_workers=1, ) ) dataloaders.append( EmmentalDataLoader( task_to_label_dict={"Thumbnail": "labels"}, dataset=test_dataset, split="test", shuffle=False, batch_size=batch_size, num_workers=1, ) ) model = EmmentalModel(name=f"Thumbnail") model.add_task( create_task("Thumbnail", n_class=2, model="resnet18", pretrained=True) ) emmental_learner = EmmentalLearner() emmental_learner.learn(model, dataloaders) scores = model.score(dataloaders) logger.warning("Model Score:") logger.warning(f"precision: {scores['Thumbnail/Thumbnail/test/precision']:.3f}") logger.warning(f"recall: {scores['Thumbnail/Thumbnail/test/recall']:.3f}") logger.warning(f"f1: {scores['Thumbnail/Thumbnail/test/f1']:.3f}")
def test_e2e(caplog): """Run an end-to-end test.""" caplog.set_level(logging.INFO) dirpath = "temp_test_e2e" Meta.reset() emmental.init(dirpath) # Generate synthetic data N = 50 X = np.random.random((N, 2)) * 2 - 1 Y1 = (X[:, 0] > X[:, 1] + 0.25).astype(int) + 1 Y2 = (-X[:, 0] > X[:, 1] + 0.25).astype(int) + 1 # Create dataset and dataloader splits = [0.8, 0.1, 0.1] X_train, X_dev, X_test = [], [], [] Y1_train, Y1_dev, Y1_test = [], [], [] Y2_train, Y2_dev, Y2_test = [], [], [] for i in range(N): if i <= N * splits[0]: X_train.append(torch.Tensor(X[i])) Y1_train.append(Y1[i]) Y2_train.append(Y2[i]) elif i < N * (splits[0] + splits[1]): X_dev.append(torch.Tensor(X[i])) Y1_dev.append(Y1[i]) Y2_dev.append(Y2[i]) else: X_test.append(torch.Tensor(X[i])) Y1_test.append(Y1[i]) Y2_test.append(Y2[i]) Y1_train = torch.from_numpy(np.array(Y1_train)) Y1_dev = torch.from_numpy(np.array(Y1_dev)) Y1_test = torch.from_numpy(np.array(Y1_test)) Y2_train = torch.from_numpy(np.array(Y1_train)) Y2_dev = torch.from_numpy(np.array(Y2_dev)) Y2_test = torch.from_numpy(np.array(Y2_test)) train_dataset1 = EmmentalDataset( name="synthetic", X_dict={"data": X_train}, Y_dict={"label1": Y1_train} ) train_dataset2 = EmmentalDataset( name="synthetic", X_dict={"data": X_train}, Y_dict={"label2": Y2_train} ) dev_dataset1 = EmmentalDataset( name="synthetic", X_dict={"data": X_dev}, Y_dict={"label1": Y1_dev} ) dev_dataset2 = EmmentalDataset( name="synthetic", X_dict={"data": X_dev}, Y_dict={"label2": Y2_dev} ) test_dataset1 = EmmentalDataset( name="synthetic", X_dict={"data": X_test}, Y_dict={"label1": Y2_test} ) test_dataset2 = EmmentalDataset( name="synthetic", X_dict={"data": X_test}, Y_dict={"label2": Y2_test} ) task_to_label_dict = {"task1": "label1"} train_dataloader1 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=train_dataset1, split="train", batch_size=10, ) dev_dataloader1 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=dev_dataset1, split="valid", batch_size=10, ) test_dataloader1 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=test_dataset1, split="test", batch_size=10, ) task_to_label_dict = {"task2": "label2"} train_dataloader2 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=train_dataset2, split="train", batch_size=10, ) dev_dataloader2 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=dev_dataset2, split="valid", batch_size=10, ) test_dataloader2 = EmmentalDataLoader( task_to_label_dict=task_to_label_dict, dataset=test_dataset2, split="test", batch_size=10, ) # Create task def ce_loss(task_name, immediate_ouput_dict, Y, active): module_name = f"{task_name}_pred_head" return F.cross_entropy( immediate_ouput_dict[module_name][0][active], (Y.view(-1) - 1)[active] ) def output(task_name, immediate_ouput_dict): module_name = f"{task_name}_pred_head" return F.softmax(immediate_ouput_dict[module_name][0], dim=1) task_name = "task1" task1 = EmmentalTask( name=task_name, module_pool=nn.ModuleDict( {"input_module": nn.Linear(2, 8), f"{task_name}_pred_head": nn.Linear(8, 2)} ), task_flow=[ { "name": "input", "module": "input_module", "inputs": [("_input_", "data")], }, { "name": f"{task_name}_pred_head", "module": f"{task_name}_pred_head", "inputs": [("input", 0)], }, ], loss_func=partial(ce_loss, task_name), output_func=partial(output, task_name), scorer=Scorer(metrics=["accuracy", "roc_auc"]), ) task_name = "task2" task2 = EmmentalTask( name=task_name, module_pool=nn.ModuleDict( {"input_module": nn.Linear(2, 8), f"{task_name}_pred_head": nn.Linear(8, 2)} ), task_flow=[ { "name": "input", "module": "input_module", "inputs": [("_input_", "data")], }, { "name": f"{task_name}_pred_head", "module": f"{task_name}_pred_head", "inputs": [("input", 0)], }, ], loss_func=partial(ce_loss, task_name), output_func=partial(output, task_name), scorer=Scorer(metrics=["accuracy", "roc_auc"]), ) # Build model mtl_model = EmmentalModel(name="all", tasks=[task1, task2]) # Create learner emmental_learner = EmmentalLearner() # Update learning config Meta.update_config( config={"learner_config": {"n_epochs": 10, "optimizer_config": {"lr": 0.01}}} ) # Learning emmental_learner.learn( mtl_model, [train_dataloader1, train_dataloader2, dev_dataloader1, dev_dataloader2], ) test1_score = mtl_model.score(test_dataloader1) test2_score = mtl_model.score(test_dataloader2) assert test1_score["task1/synthetic/test/accuracy"] >= 0.5 assert test1_score["task1/synthetic/test/roc_auc"] >= 0.6 assert test2_score["task2/synthetic/test/accuracy"] >= 0.5 assert test2_score["task2/synthetic/test/roc_auc"] >= 0.6 shutil.rmtree(dirpath)
def run_model(mode, config, run_config_path=None): """ Main run method for Emmental Bootleg models. Args: mode: run mode (train, eval, dump_preds, dump_embs) config: parsed model config run_config_path: original config path (for saving) Returns: """ # Set up distributed backend and save configuration files setup(config, run_config_path) # Load entity symbols log_rank_0_info(logger, f"Loading entity symbols...") entity_symbols = EntitySymbols.load_from_cache( load_dir=os.path.join(config.data_config.entity_dir, config.data_config.entity_map_dir), alias_cand_map_file=config.data_config.alias_cand_map, alias_idx_file=config.data_config.alias_idx_map, ) # Create tasks tasks = [NED_TASK] if config.data_config.type_prediction.use_type_pred is True: tasks.append(TYPE_PRED_TASK) # Create splits for data loaders data_splits = [TRAIN_SPLIT, DEV_SPLIT, TEST_SPLIT] # Slices are for eval so we only split on test/dev slice_splits = [DEV_SPLIT, TEST_SPLIT] # If doing eval, only run on test data if mode in ["eval", "dump_preds", "dump_embs"]: data_splits = [TEST_SPLIT] slice_splits = [TEST_SPLIT] # We only do dumping if weak labels is True if mode in ["dump_preds", "dump_embs"]: if config.data_config[ f"{TEST_SPLIT}_dataset"].use_weak_label is False: raise ValueError( f"When calling dump_preds or dump_embs, we require use_weak_label to be True." ) # Gets embeddings that need to be prepped during data prep or in the __get_item__ method batch_on_the_fly_kg_adj = get_dataloader_embeddings(config, entity_symbols) # Gets dataloaders dataloaders = get_dataloaders( config, tasks, data_splits, entity_symbols, batch_on_the_fly_kg_adj, ) slice_datasets = get_slicedatasets(config, slice_splits, entity_symbols) configure_optimizer(config) # Create models and add tasks if config.model_config.attn_class == "BERTNED": log_rank_0_info(logger, f"Starting NED-Base Model") assert (config.data_config.type_prediction.use_type_pred is False), f"NED-Base does not support type prediction" assert ( config.data_config.word_embedding.use_sent_proj is False ), f"NED-Base requires word_embeddings.use_sent_proj to be False" model = EmmentalModel(name="NED-Base") model.add_tasks( ned_task.create_task(config, entity_symbols, slice_datasets)) else: log_rank_0_info(logger, f"Starting Bootleg Model") model = EmmentalModel(name="Bootleg") # TODO: make this more general for other tasks -- iterate through list of tasks # and add task for each model.add_task( ned_task.create_task(config, entity_symbols, slice_datasets)) if TYPE_PRED_TASK in tasks: model.add_task( type_pred_task.create_task(config, entity_symbols, slice_datasets)) # Add the mention type embedding to the embedding payload type_pred_task.update_ned_task(model) # Print param counts if mode == "train": log_rank_0_debug(logger, "PARAMS WITH GRAD\n" + "=" * 30) total_params = count_parameters(model, requires_grad=True, logger=logger) log_rank_0_info(logger, f"===> Total Params With Grad: {total_params}") log_rank_0_debug(logger, "PARAMS WITHOUT GRAD\n" + "=" * 30) total_params = count_parameters(model, requires_grad=False, logger=logger) log_rank_0_info(logger, f"===> Total Params Without Grad: {total_params}") # Load the best model from the pretrained model if config["model_config"]["model_path"] is not None: model.load(config["model_config"]["model_path"]) # Barrier if config["learner_config"]["local_rank"] == 0: torch.distributed.barrier() # Train model if mode == "train": emmental_learner = EmmentalLearner() emmental_learner._set_optimizer(model) emmental_learner.learn(model, dataloaders) if config.learner_config.local_rank in [0, -1]: model.save(f"{emmental.Meta.log_path}/last_model.pth") # Multi-gpu DataParallel eval (NOT distributed) if mode in ["eval", "dump_embs", "dump_preds"]: # This happens inside EmmentalLearner for training if (config["learner_config"]["local_rank"] == -1 and config["model_config"]["dataparallel"]): model._to_dataparallel() # If just finished training a model or in eval mode, run eval if mode in ["train", "eval"]: scores = model.score(dataloaders) # Save metrics and models log_rank_0_info(logger, f"Saving metrics to {emmental.Meta.log_path}") log_rank_0_info(logger, f"Metrics: {scores}") scores["log_path"] = emmental.Meta.log_path if config.learner_config.local_rank in [0, -1]: write_to_file(f"{emmental.Meta.log_path}/{mode}_metrics.txt", scores) eval_utils.write_disambig_metrics_to_csv( f"{emmental.Meta.log_path}/{mode}_disambig_metrics.csv", scores) return scores # If you want detailed dumps, save model outputs assert mode in [ "dump_preds", "dump_embs", ], 'Mode must be "dump_preds" or "dump_embs"' dump_embs = False if mode != "dump_embs" else True assert ( len(dataloaders) == 1 ), f"We should only have length 1 dataloaders for dump_embs and dump_preds!" final_result_file, final_out_emb_file = None, None if config.learner_config.local_rank in [0, -1]: # Setup files/folders filename = os.path.basename(dataloaders[0].dataset.raw_filename) log_rank_0_debug( logger, f"Collecting sentence to mention map {os.path.join(config.data_config.data_dir, filename)}", ) sentidx2num_mentions, sent_idx2row = eval_utils.get_sent_idx2num_mens( os.path.join(config.data_config.data_dir, filename)) log_rank_0_debug(logger, f"Done collecting sentence to mention map") eval_folder = eval_utils.get_eval_folder(filename) subeval_folder = os.path.join(eval_folder, "batch_results") utils.ensure_dir(subeval_folder) # Will keep track of sentences dumped already. These will only be ones with mentions all_dumped_sentences = set() number_dumped_batches = 0 total_mentions_seen = 0 all_result_files = [] all_out_emb_files = [] # Iterating over batches of predictions for res_i, res_dict in enumerate( eval_utils.batched_pred_iter( model, dataloaders[0], config.run_config.eval_accumulation_steps, sentidx2num_mentions, )): ( result_file, out_emb_file, final_sent_idxs, mentions_seen, ) = eval_utils.disambig_dump_preds( res_i, total_mentions_seen, config, res_dict, sentidx2num_mentions, sent_idx2row, subeval_folder, entity_symbols, dump_embs, NED_TASK, ) all_dumped_sentences.update(final_sent_idxs) all_result_files.append(result_file) all_out_emb_files.append(out_emb_file) total_mentions_seen += mentions_seen number_dumped_batches += 1 # Dump the sentences that had no mentions and were not already dumped # Assert all remaining sentences have no mentions assert all( v == 0 for k, v in sentidx2num_mentions.items() if k not in all_dumped_sentences ), (f"Sentences with mentions were not dumped: " f"{[k for k, v in sentidx2num_mentions.items() if k not in all_dumped_sentences]}" ) empty_sentidx2row = { k: v for k, v in sent_idx2row.items() if k not in all_dumped_sentences } empty_resultfile = eval_utils.get_result_file(number_dumped_batches, subeval_folder) all_result_files.append(empty_resultfile) # Dump the outputs eval_utils.write_data_labels_single( sentidx2row=empty_sentidx2row, output_file=empty_resultfile, filt_emb_data=None, sental2embid={}, alias_cand_map=entity_symbols.get_alias2qids(), qid2eid=entity_symbols.get_qid2eid(), result_alias_offset=total_mentions_seen, train_in_cands=config.data_config.train_in_candidates, max_cands=entity_symbols.max_candidates, dump_embs=dump_embs, ) log_rank_0_info( logger, f"Finished dumping. Merging results across accumulation steps.") # Final result files for labels and embeddings final_result_file = os.path.join(eval_folder, config.run_config.result_label_file) # Copy labels output = open(final_result_file, "wb") for file in all_result_files: shutil.copyfileobj(open(file, "rb"), output) output.close() log_rank_0_info(logger, f"Bootleg labels saved at {final_result_file}") # Try to copy embeddings if dump_embs: final_out_emb_file = os.path.join( eval_folder, config.run_config.result_emb_file) log_rank_0_info( logger, f"Trying to merge numpy embedding arrays. " f"If your machine is limited in memory, this may cause OOM errors. " f"Is that happens, result files should be saved in {subeval_folder}.", ) all_arrays = [] for i, npfile in enumerate(all_out_emb_files): all_arrays.append(np.load(npfile)) np.save(final_out_emb_file, np.concatenate(all_arrays)) log_rank_0_info( logger, f"Bootleg embeddings saved at {final_out_emb_file}") # Cleanup try_rmtree(subeval_folder) return final_result_file, final_out_emb_file