def test_asgd_optimizer(caplog): """Unit test of ASGD optimizer""" caplog.set_level(logging.INFO) optimizer = "asgd" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default ASGD setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert isequal( emmental_learner.optimizer.defaults, { "lr": 0.001, "lambd": 0.0001, "alpha": 0.75, "t0": 1_000_000.0, "weight_decay": 0, }, )
def test_lbfgs_optimizer(caplog): """Unit test of LBFGS optimizer""" caplog.set_level(logging.INFO) optimizer = "lbfgs" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default LBFGS setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "max_iter": 20, "max_eval": 25, "tolerance_grad": 1e-07, "tolerance_change": 1e-09, "history_size": 100, "line_search_fn": None, } # Test new LBFGS setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "max_iter": 30, "max_eval": 40, "tolerance_grad": 1e-04, "tolerance_change": 1e-05, "history_size": 10, "line_search_fn": "strong_wolfe", }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "max_iter": 30, "max_eval": 40, "tolerance_grad": 1e-04, "tolerance_change": 1e-05, "history_size": 10, "line_search_fn": "strong_wolfe", } shutil.rmtree(dirpath)
def test_adagrad_optimizer(caplog): """Unit test of Adagrad optimizer""" caplog.set_level(logging.INFO) optimizer = "adagrad" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default Adagrad setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert isequal( emmental_learner.optimizer.defaults, { "lr": 0.001, "lr_decay": 0, "initial_accumulator_value": 0, "eps": 1e-10, "weight_decay": 0, }, ) # Test new Adagrad setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "lr_decay": 0.1, "initial_accumulator_value": 0.2, "eps": 1e-5, }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert isequal( emmental_learner.optimizer.defaults, { "lr": 0.02, "lr_decay": 0.1, "initial_accumulator_value": 0.2, "eps": 1e-5, "weight_decay": 0.05, }, ) shutil.rmtree(dirpath)
def test_adadelta_optimizer(caplog): """Unit test of Adadelta optimizer""" caplog.set_level(logging.INFO) optimizer = "adadelta" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default Adadelta setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert isequal( emmental_learner.optimizer.defaults, { "lr": 0.001, "rho": 0.9, "eps": 1e-06, "weight_decay": 0 }, ) # Test new Adadelta setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "rho": 0.6, "eps": 1e-05 }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert isequal( emmental_learner.optimizer.defaults, { "lr": 0.02, "rho": 0.6, "eps": 1e-05, "weight_decay": 0.05 }, ) shutil.rmtree(dirpath)
def test_plateau_scheduler(caplog): """Unit test of plateau scheduler.""" caplog.set_level(logging.INFO) lr_scheduler = "plateau" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) config = { "learner_config": { "n_epochs": 4, "optimizer_config": {"optimizer": "sgd", "lr": 10}, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "plateau_config": { "metric": "model/train/all/loss", "mode": "min", "factor": 0.1, "patience": 1, "threshold": 0.0001, "threshold_mode": "rel", "cooldown": 0, "eps": 1e-08, }, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_learning_counter() emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {"model/train/all/loss": 1}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {"model/train/all/loss": 1}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {"model/train/all/loss": 1}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {"model/train/all/loss": 0.1}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 shutil.rmtree(dirpath)
def test_bert_adam_optimizer(caplog): """Unit test of BertAdam optimizer""" caplog.set_level(logging.INFO) optimizer = "bert_adam" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default BertAdam setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "betas": (0.9, 0.999), "eps": 1e-08, "weight_decay": 0.0, } # Test new BertAdam setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "betas": (0.8, 0.9), "eps": 1e-05 }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "betas": (0.8, 0.9), "eps": 1e-05, "weight_decay": 0.05, } # Test BertAdam setp emmental_learner.optimizer.zero_grad() torch.Tensor(1) F.mse_loss(model(torch.randn(1, 1)), torch.randn(1, 1)).backward() emmental_learner.optimizer.step() shutil.rmtree(dirpath)
def test_rms_prop_optimizer(caplog): """Unit test of RMSprop optimizer""" caplog.set_level(logging.INFO) optimizer = "rms_prop" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default RMSprop setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "alpha": 0.99, "eps": 1e-08, "momentum": 0, "centered": False, "weight_decay": 0, } # Test new RMSprop setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "alpha": 0.9, "eps": 1e-05, "momentum": 0.1, "centered": True, }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "alpha": 0.9, "eps": 1e-05, "momentum": 0.1, "centered": True, "weight_decay": 0.05, } shutil.rmtree(dirpath)
def test_step_scheduler(caplog): """Unit test of step scheduler.""" caplog.set_level(logging.INFO) lr_scheduler = "step" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) config = { "learner_config": { "n_epochs": 4, "optimizer_config": { "optimizer": "sgd", "lr": 10 }, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "step_config": { "step_size": 2, "gamma": 0.1, "last_epoch": -1 }, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_learning_counter() emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.1) < 1e-5 shutil.rmtree(dirpath)
def test_adamw_optimizer(caplog): """Unit test of AdamW optimizer""" caplog.set_level(logging.INFO) optimizer = "adamw" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default AdamW setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "betas": (0.9, 0.999), "eps": 1e-08, "amsgrad": False, "weight_decay": 0, } # Test new AdamW setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "betas": (0.9, 0.99), "eps": 1e-05, "amsgrad": True, }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "betas": (0.9, 0.99), "eps": 1e-05, "amsgrad": True, "weight_decay": 0.05, } shutil.rmtree(dirpath)
def test_sgd_optimizer(caplog): """Unit test of SGD optimizer.""" caplog.set_level(logging.INFO) optimizer = "sgd" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default SGD setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "momentum": 0, "dampening": 0, "nesterov": False, "weight_decay": 0.0, } # Test new SGD setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "momentum": 0.1, "dampening": 0, "nesterov": True, }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "momentum": 0.1, "dampening": 0, "nesterov": True, "weight_decay": 0.05, } shutil.rmtree(dirpath)
def test_cosine_annealing_scheduler(caplog): """Unit test of cosine annealing scheduler""" caplog.set_level(logging.INFO) lr_scheduler = "cosine_annealing" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) config = { "learner_config": { "n_epochs": 4, "optimizer_config": { "optimizer": "sgd", "lr": 10 }, "lr_scheduler_config": { "lr_scheduler": lr_scheduler }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert (abs(emmental_learner.optimizer.param_groups[0]["lr"] - 8.535533905932738) < 1e-5) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert (abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1.4644660940672627) < 1e-5) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"]) < 1e-5 shutil.rmtree(dirpath)
def test_sparse_adam_optimizer(caplog): """Unit test of SparseAdam optimizer""" caplog.set_level(logging.INFO) optimizer = "sparse_adam" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default SparseAdam setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "betas": (0.9, 0.999), "eps": 1e-08, } # Test new SparseAdam setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "betas": (0.8, 0.9), "eps": 1e-05 }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "betas": (0.8, 0.9), "eps": 1e-05, } shutil.rmtree(dirpath)
def test_r_prop_optimizer(caplog): """Unit test of Rprop optimizer""" caplog.set_level(logging.INFO) optimizer = "r_prop" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test default Rprop setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "etas": (0.5, 1.2), "step_sizes": (1e-06, 50), } # Test new Rprop setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "etas": (0.3, 1.5), "step_sizes": (1e-04, 30) }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "etas": (0.3, 1.5), "step_sizes": (1e-04, 30), } shutil.rmtree(dirpath)
def test_linear_scheduler(caplog): """Unit test of linear scheduler.""" caplog.set_level(logging.INFO) lr_scheduler = "linear" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test per batch config = { "learner_config": { "n_epochs": 4, "optimizer_config": { "optimizer": "sgd", "lr": 10 }, "lr_scheduler_config": { "lr_scheduler": lr_scheduler }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_learning_counter() emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 7.5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 2.5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"]) < 1e-5 # Test every 2 batches config = { "learner_config": { "n_epochs": 4, "optimizer_config": { "optimizer": "sgd", "lr": 10 }, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "lr_scheduler_step_freq": 2, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 7.5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 7.5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 5.0) < 1e-5 shutil.rmtree(dirpath)
def test_one_cycle_scheduler(caplog): """Unit test of one cycle scheduler""" caplog.set_level(logging.INFO) lr_scheduler = "one_cycle" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) config = { "learner_config": { "n_epochs": 4, "optimizer_config": {"optimizer": "sgd", "lr": 10}, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "one_cycle_config": { "anneal_strategy": "cos", "base_momentum": 0.85, "cycle_momentum": True, "div_factor": 1, "final_div_factor": 10000.0, "last_epoch": -1, "max_lr": 0.1, "max_momentum": 0.95, "pct_start": 0.3, }, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 0.1 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert ( abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.08117637264392738) < 1e-5 ) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert ( abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.028312982462817687) < 1e-5 ) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1e-05) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert ( abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.028312982462817677) < 1e-5 ) 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
def test_sparse_adam_optimizer(caplog): """Unit test of SparseAdam optimizer.""" caplog.set_level(logging.INFO) optimizer = "sparse_adam" dirpath = "temp_test_optimizer" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) def grouped_parameters(model): no_decay = ["bias"] 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 # Test default SparseAdam setting config = {"learner_config": {"optimizer_config": {"optimizer": optimizer}}} emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.001, "betas": (0.9, 0.999), "eps": 1e-08, } # Test new SparseAdam setting config = { "learner_config": { "optimizer_config": { "optimizer": optimizer, "lr": 0.02, "l2": 0.05, f"{optimizer}_config": { "betas": (0.8, 0.9), "eps": 1e-05 }, } } } emmental.Meta.update_config(config) emmental_learner._set_optimizer(model) assert emmental_learner.optimizer.defaults == { "lr": 0.02, "betas": (0.8, 0.9), "eps": 1e-05, } shutil.rmtree(dirpath)
def test_exponential_scheduler(caplog): """Unit test of exponential scheduler""" caplog.set_level(logging.INFO) lr_scheduler = "exponential" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test step per batch config = { "learner_config": { "n_epochs": 4, "optimizer_config": {"optimizer": "sgd", "lr": 10}, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "exponential_config": {"gamma": 0.1}, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.01) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.001) < 1e-5 # Test step per epoch config = { "learner_config": { "n_epochs": 4, "optimizer_config": {"optimizer": "sgd", "lr": 10}, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "lr_scheduler_step_unit": "epoch", "exponential_config": {"gamma": 0.1}, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 2 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 1) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 0.1) < 1e-5 shutil.rmtree(dirpath)
def test_step_scheduler(caplog): """Unit test of step scheduler.""" caplog.set_level(logging.INFO) dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) # Test warmup steps config = { "learner_config": { "n_epochs": 4, "optimizer_config": {"optimizer": "sgd", "lr": 10}, "lr_scheduler_config": { "lr_scheduler": None, "warmup_steps": 2, "warmup_unit": "batch", }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 0 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 Meta.reset() emmental.init(dirpath) # Test warmup percentage config = { "learner_config": { "n_epochs": 4, "optimizer_config": {"optimizer": "sgd", "lr": 10}, "lr_scheduler_config": { "lr_scheduler": None, "warmup_percentage": 0.5, "warmup_unit": "epoch", }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 0 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 5) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert abs(emmental_learner.optimizer.param_groups[0]["lr"] - 10) < 1e-5 shutil.rmtree(dirpath)
def test_cyclic_scheduler(caplog): """Unit test of cyclic scheduler""" caplog.set_level(logging.INFO) lr_scheduler = "cyclic" dirpath = "temp_test_scheduler" model = nn.Linear(1, 1) emmental_learner = EmmentalLearner() Meta.reset() emmental.init(dirpath) config = { "learner_config": { "n_epochs": 4, "optimizer_config": { "optimizer": "sgd", "lr": 10 }, "lr_scheduler_config": { "lr_scheduler": lr_scheduler, "cyclic_config": { "base_lr": 10, "base_momentum": 0.8, "cycle_momentum": True, "gamma": 1.0, "last_epoch": -1, "max_lr": 0.1, "max_momentum": 0.9, "mode": "triangular", "scale_fn": None, "scale_mode": "cycle", "step_size_down": None, "step_size_up": 2000, }, }, } } emmental.Meta.update_config(config) emmental_learner.n_batches_per_epoch = 1 emmental_learner._set_optimizer(model) emmental_learner._set_lr_scheduler(model) assert emmental_learner.optimizer.param_groups[0]["lr"] == 10 emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 0, {}) assert (abs(emmental_learner.optimizer.param_groups[0]["lr"] - 9.995049999999999) < 1e-5) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 1, {}) assert (abs(emmental_learner.optimizer.param_groups[0]["lr"] - 9.990100000000002) < 1e-5) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 2, {}) assert (abs(emmental_learner.optimizer.param_groups[0]["lr"] - 9.985149999999999) < 1e-5) emmental_learner.optimizer.step() emmental_learner._update_lr_scheduler(model, 3, {}) assert (abs(emmental_learner.optimizer.param_groups[0]["lr"] - 9.980200000000002) < 1e-5) shutil.rmtree(dirpath)