def test_tensorboard_writer(caplog):
"""Unit test of log_writer."""
caplog.set_level(logging.INFO)
emmental.Meta.reset()
emmental.init()
log_writer = TensorBoardWriter()
log_writer.add_scalar(name="step 1", value=0.1, step=1)
log_writer.add_scalar(name="step 2", value=0.2, step=2)
config_filename = "config.yaml"
log_writer.write_config(config_filename)
# Test config
with open(os.path.join(emmental.Meta.log_path, config_filename), "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)
assert config["meta_config"]["verbose"] is True
assert config["logging_config"]["counter_unit"] == "epoch"
assert config["logging_config"]["checkpointing"] is False
log_writer.write_log()
log_writer.close()
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_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_mixed_scheduler_no_y_dict(caplog):
"""Unit test of mixed scheduler with no y_dict."""
caplog.set_level(logging.INFO)
init()
task1 = "task1"
x1 = np.random.rand(20, 2)
task2 = "task2"
x2 = np.random.rand(30, 3)
dataloaders = [
EmmentalDataLoader(
task_to_label_dict={task_name: None},
dataset=EmmentalDataset(name=task_name, X_dict={"feature": x}),
split="train",
batch_size=10,
shuffle=True,
) for task_name, x in [(task1, x1), (task2, x2)]
]
scheduler = MixedScheduler()
assert scheduler.get_num_batches(dataloaders) == 2
batch_y_dict_1 = [
batch[0].Y_dict for batch in scheduler.get_batches(dataloaders)
]
batch_y_dict_2 = [
batch[1].Y_dict for batch in scheduler.get_batches(dataloaders)
]
assert batch_y_dict_1 == [None] * 2
assert batch_y_dict_2 == [None] * 2
def test_logging_manager_tensorboard(caplog):
"""Unit test of logging_manager (tensorboard)."""
caplog.set_level(logging.INFO)
emmental.init()
Meta.update_config(
config={
"logging_config": {
"counter_unit": "epoch",
"evaluation_freq": 1,
"checkpointing": False,
"checkpointer_config": {
"checkpoint_freq": 2
},
"writer_config": {
"writer": "tensorboard"
},
}
})
logging_manager = LoggingManager(n_batches_per_epoch=2)
logging_manager.update(5)
assert type(logging_manager.writer) == TensorBoardWriter
def test_checkpointer_wrong_metric_mode(caplog):
"""Unit test of checkpointer (wrong metric mode)"""
caplog.set_level(logging.INFO)
emmental.Meta.reset()
emmental.init()
emmental.Meta.update_config(
config={
"logging_config": {
"counter_unit": "sample",
"evaluation_freq": 10,
"checkpointing": True,
"checkpointer_config": {
"checkpoint_metric": {
"model/all/train/loss": "min1"
}
},
}
})
with pytest.raises(ValueError):
checkpointer = Checkpointer()
checkpointer.clear()
def test_checkpointer_metric(caplog):
"""Unit test of checkpointer (metric)"""
caplog.set_level(logging.INFO)
checkpoint_path = "temp_test_checkpointer"
emmental.Meta.reset()
emmental.init()
emmental.Meta.update_config(
config={
"logging_config": {
"counter_unit": "sample",
"evaluation_freq": 10,
"checkpointing": True,
"checkpointer_config": {
"checkpoint_metric": None,
"checkpoint_task_metrics": {
"model/all/train/loss": "min",
"model/all/train/accuracy": "max",
},
"checkpoint_freq": 2,
"checkpoint_path": checkpoint_path,
},
}
})
checkpointer = Checkpointer()
checkpointer.clear()
assert os.path.exists(checkpoint_path) is True
shutil.rmtree(checkpoint_path)
def test_checkpointer_specific_path(caplog):
"""Unit test of checkpointer"""
caplog.set_level(logging.INFO)
checkpoint_path = "temp_test_checkpointer"
emmental.Meta.reset()
emmental.init()
emmental.Meta.update_config(
config={
"logging_config": {
"counter_unit": "sample",
"evaluation_freq": 10,
"checkpointing": True,
"checkpointer_config": {
"checkpoint_freq": 2,
"checkpoint_path": checkpoint_path,
},
}
})
checkpointer = Checkpointer()
checkpointer.clear()
assert os.path.exists(checkpoint_path) is True
shutil.rmtree(checkpoint_path)
def setUp(self) -> None:
self.args = parser_utils.parse_boot_and_emm_args(
"test/run_args/test_end2end.json")
# This _MUST_ get passed the args so it gets a random seed set
emmental.init(log_dir="test/temp_log", config=self.args)
if not os.path.exists(emmental.Meta.log_path):
os.makedirs(emmental.Meta.log_path)
def test_logging_manager_no_writer(caplog):
"""Unit test of logging_manager (no writer)."""
caplog.set_level(logging.INFO)
emmental.init()
Meta.update_config(
config={
"logging_config": {
"counter_unit": "epoch",
"evaluation_freq": 1,
"checkpointing": False,
"checkpointer_config": {
"checkpoint_freq": 2
},
"writer_config": {
"writer": None
},
}
})
logging_manager = LoggingManager(n_batches_per_epoch=2)
logging_manager.update(5)
assert logging_manager.writer is None
def setup(config, run_config_path=None):
"""
Setup distributed backend and save configuration files.
Args:
config: config
run_config_path: path for original run config
Returns:
"""
# torch.multiprocessing.set_sharing_strategy("file_system")
# spawn method must be fork to work with Meta.config
torch.multiprocessing.set_start_method("fork", force=True)
"""
ulimit -n 500000
python3 -m torch.distributed.launch --nproc_per_node=2 bootleg/run.py --config_script ...
"""
log_level = logging.getLevelName(config.run_config.log_level.upper())
emmental.init(
log_dir=config["meta_config"]["log_path"],
config=config,
use_exact_log_path=config["meta_config"]["use_exact_log_path"],
local_rank=config.learner_config.local_rank,
level=log_level,
)
log = logging.getLogger()
# Remove streaming handlers and use rich
log.handlers = [
h for h in log.handlers if not type(h) is logging.StreamHandler
]
log.addHandler(RichHandler())
# Set up distributed backend
emmental.Meta.init_distributed_backend()
cmd_msg = " ".join(sys.argv)
# Log configuration into filess
if config.learner_config.local_rank in [0, -1]:
write_to_file(f"{emmental.Meta.log_path}/cmd.txt", cmd_msg)
dump_yaml_file(f"{emmental.Meta.log_path}/parsed_config.yaml",
emmental.Meta.config)
# Dump the run config (does not contain defaults)
if run_config_path is not None:
dump_yaml_file(
f"{emmental.Meta.log_path}/run_config.yaml",
load_yaml_file(run_config_path),
)
log_rank_0_info(logger, f"COMMAND: {cmd_msg}")
log_rank_0_info(
logger,
f"Saving config to {emmental.Meta.log_path}/parsed_config.yaml")
git_hash = "Not able to retrieve git hash"
try:
git_hash = subprocess.check_output([
"git", "log", "-n", "1", "--pretty=tformat:%h-%ad", "--date=short"
]).strip()
except subprocess.CalledProcessError:
pass
log_rank_0_info(logger, f"Git Hash: {git_hash}")
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_config_check_in_meta(caplog):
"""Unit test of config check in meta."""
caplog.set_level(logging.INFO)
dirpath = "temp_test_meta_log_folder"
Meta.reset()
init(dirpath)
config = {
"logging_config": {
"evaluation_freq": 5.0,
},
}
Meta.update_config(config)
assert type(Meta.config["logging_config"]["evaluation_freq"]) == int
assert Meta.config["logging_config"]["evaluation_freq"] == 5
config = {
"logging_config": {
"counter_unit": "batch",
"evaluation_freq": 2.3,
},
}
Meta.update_config(config)
assert type(Meta.config["logging_config"]["evaluation_freq"]) == int
assert Meta.config["logging_config"]["evaluation_freq"] == 3
config = {
"logging_config": {
"counter_unit": "sample",
"evaluation_freq": 0.2,
},
}
Meta.update_config(config)
assert type(Meta.config["logging_config"]["evaluation_freq"]) == int
assert Meta.config["logging_config"]["evaluation_freq"] == 1
config = {
"logging_config": {
"counter_unit": "epoch",
"evaluation_freq": 1,
"writer_config": {
"write_loss_per_step": True
},
},
}
Meta.update_config(config)
assert (
Meta.config["logging_config"]["writer_config"]["write_loss_per_step"]
is False)
# Remove the temp folder
shutil.rmtree(dirpath)
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 setUp(self) -> None:
emmental.init(log_dir="test/temp_log")
if not os.path.exists(emmental.Meta.log_path):
os.makedirs(emmental.Meta.log_path)
self.args = parser_utils.parse_boot_and_emm_args(
"test/run_args/test_embeddings.json"
)
self.entity_symbols = EntitySymbolsSubclass()
self.tokenizer = load_tokenizer()
def _init_meta(self, _seed, exp_dir, meta_config, learner_config, logging_config):
emmental.init(path.join(exp_dir, '_emmental_logs'))
Meta.update_config(
config={
'meta_config': {**meta_config, 'seed': _seed},
'learner_config': learner_config,
'logging_config': logging_config
}
)
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_emmental_task(caplog):
"""Unit test of emmental task."""
caplog.set_level(logging.INFO)
emmental.init()
def ce_loss(module_name, output_dict, Y):
return F.cross_entropy(output_dict[module_name][0], Y.view(-1))
def output(module_name, output_dict):
return F.softmax(output_dict[module_name][0], dim=1)
task_name = "task1"
task_metrics = {task_name: ["accuracy"]}
scorer = Scorer(metrics=task_metrics[task_name])
task = EmmentalTask(
name=task_name,
module_pool=nn.ModuleDict({
"input_module0": IdentityModule(),
"input_module1": IdentityModule(),
f"{task_name}_pred_head": IdentityModule(),
}),
task_flow=[
Action("input1", "input_module0", [("_input_", "data")]),
Action("input2", "input_module1", [("input1", 0)]),
Action(f"{task_name}_pred_head", f"{task_name}_pred_head",
[("input2", 0)]),
],
module_device={
"input_module0": -1,
"input_module1": 0,
"input_module": -1
},
loss_func=partial(ce_loss, f"{task_name}_pred_head"),
output_func=partial(output, f"{task_name}_pred_head"),
action_outputs=None,
scorer=scorer,
require_prob_for_eval=False,
require_pred_for_eval=True,
weight=2.0,
)
assert task.name == task_name
assert set(list(task.module_pool.keys())) == set(
["input_module0", "input_module1", f"{task_name}_pred_head"])
assert task.action_outputs is None
assert task.scorer == scorer
assert len(task.task_flow) == 3
assert task.module_device == {
"input_module0": torch.device("cpu"),
"input_module1": torch.device(0),
}
assert task.require_prob_for_eval is False
assert task.require_pred_for_eval is True
assert task.weight == 2.0
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_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_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 setUp(self) -> None:
emmental.init(log_dir="test/temp_log")
if not os.path.exists(emmental.Meta.log_path):
os.makedirs(emmental.Meta.log_path)
self.args = parser_utils.parse_boot_and_emm_args(
"test/run_args/test_embeddings.json"
)
self.entity_symbols = EntitySymbolsSubclass()
self.kg_adj = os.path.join(self.args.data_config.emb_dir, "kg_adj.txt")
self.kg_adj_json = os.path.join(self.args.data_config.emb_dir, "kg_adj.json")
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_mixed_scheduler(caplog):
"""Unit test of mixed scheduler."""
caplog.set_level(logging.INFO)
init()
task1 = "task1"
x1 = np.random.rand(20, 2)
y1 = torch.from_numpy(np.random.rand(20))
task2 = "task2"
x2 = np.random.rand(30, 3)
y2 = torch.from_numpy(np.random.rand(30))
dataloaders = [
EmmentalDataLoader(
task_to_label_dict={task_name: "label"},
dataset=EmmentalDataset(
name=task_name, X_dict={"feature": x}, Y_dict={"label": y}
),
split="train",
batch_size=10,
shuffle=True,
)
for task_name, x, y in [(task1, x1, y1), (task2, x2, y2)]
]
scheduler = MixedScheduler()
assert scheduler.get_num_batches(dataloaders) == 2
batch_task_names_1 = [
batch_data[0][-2] for batch_data in scheduler.get_batches(dataloaders)
]
batch_task_names_2 = [
batch_data[1][-2] for batch_data in scheduler.get_batches(dataloaders)
]
assert batch_task_names_1 == [task1, task1]
assert batch_task_names_2 == [task2, task2]
scheduler = MixedScheduler(fillup=True)
assert scheduler.get_num_batches(dataloaders) == 3
batch_task_names_1 = [
batch_data[0][-2] for batch_data in scheduler.get_batches(dataloaders)
]
batch_task_names_2 = [
batch_data[1][-2] for batch_data in scheduler.get_batches(dataloaders)
]
assert batch_task_names_1 == [task1, task1, task1]
assert batch_task_names_2 == [task2, task2, task2]
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 _init_meta(self, seed, exp_dir, meta_config, learner_config, logging_config):
emmental.init(path.join(exp_dir, '_emmental_logs'))
print(seed)
Meta.update_config(
config={
'meta_config': {**meta_config, 'seed': seed},
'model_config': {'device': meta_config['device']},
'learner_config': learner_config,
'logging_config': logging_config
}
)
print(Meta.config)
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 setUp(self) -> None:
self.args = parser_utils.parse_boot_and_emm_args(
"test/run_args/test_embeddings.json"
)
emmental.init(log_dir="test/temp_log", config=self.args)
if not os.path.exists(emmental.Meta.log_path):
os.makedirs(emmental.Meta.log_path)
self.args.data_config.ent_embeddings = [
DottedDict(
{
"key": "learned1",
"load_class": "LearnedEntityEmb",
"dropout1d": 0.5,
"args": {"learned_embedding_size": 5, "tail_init": False},
}
),
DottedDict(
{
"key": "learned2",
"dropout2d": 0.5,
"load_class": "LearnedEntityEmb",
"args": {"learned_embedding_size": 5, "tail_init": False},
}
),
DottedDict(
{
"key": "learned3",
"load_class": "LearnedEntityEmb",
"freeze": True,
"args": {"learned_embedding_size": 5, "tail_init": False},
}
),
DottedDict(
{
"key": "learned4",
"load_class": "LearnedEntityEmb",
"normalize": False,
"args": {"learned_embedding_size": 5, "tail_init": False},
}
),
DottedDict(
{
"key": "learned5",
"load_class": "LearnedEntityEmb",
"cpu": True,
"args": {"learned_embedding_size": 5, "tail_init": False},
}
),
]
self.tokenizer = load_tokenizer()
self.entity_symbols = EntitySymbolsSubclass()
def test_logging_manager_no_checkpointing(caplog):
"""Unit test of logging_manager (no checkpointing)"""
caplog.set_level(logging.INFO)
emmental.init()
Meta.update_config(
config={
"meta_config": {
"verbose": False
},
"logging_config": {
"counter_unit": "epoch",
"evaluation_freq": 1,
"checkpointing": False,
"checkpointer_config": {
"checkpoint_freq": 2
},
"writer_config": {
"writer": "json"
},
},
})
logging_manager = LoggingManager(n_batches_per_epoch=2)
logging_manager.update(5)
assert logging_manager.trigger_evaluation() is False
assert logging_manager.trigger_checkpointing() is False
logging_manager.update(5)
assert logging_manager.trigger_evaluation() is True
assert logging_manager.trigger_checkpointing() is False
logging_manager.update(10)
assert logging_manager.trigger_evaluation() is False
assert logging_manager.trigger_checkpointing() is False
logging_manager.update(5)
assert logging_manager.trigger_evaluation() is True
assert logging_manager.trigger_checkpointing() is False
assert logging_manager.epoch_count == 0
assert logging_manager.sample_total == 25
assert logging_manager.batch_total == 4
assert logging_manager.epoch_total == 2
model = EmmentalModel()
logging_manager.close(model)
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)
