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)
