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_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_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_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 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_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)
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)