def test_wandb_logger(wandb):
"""Verify that basic functionality of wandb logger works.
Wandb doesn't work well with pytest so we have to mock it out here."""
logger = WandbLogger(anonymous=True, offline=True)
logger.log_metrics({'acc': 1.0})
wandb.init().log.assert_called_once_with({'acc': 1.0}, step=None)
wandb.init().log.reset_mock()
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init().log.assert_called_once_with({'acc': 1.0}, step=3)
# continue training on same W&B run
wandb.init().step = 3
logger.finalize('success')
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init().log.assert_called_with({'acc': 1.0}, step=6)
logger.log_hyperparams({'test': None, 'nested': {'a': 1}, 'b': [2, 3, 4]})
wandb.init().config.update.assert_called_once_with(
{
'test': 'None',
'nested/a': 1,
'b': [2, 3, 4]
},
allow_val_change=True,
)
logger.watch('model', 'log', 10)
wandb.init().watch.assert_called_once_with('model', log='log', log_freq=10)
assert logger.name == wandb.init().project_name()
assert logger.version == wandb.init().id
def test_wandb_logger_init(wandb, recwarn):
"""Verify that basic functionality of wandb logger works.
Wandb doesn't work well with pytest so we have to mock it out here."""
# test wandb.init called when there is no W&B run
wandb.run = None
logger = WandbLogger()
logger.log_metrics({'acc': 1.0})
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({'acc': 1.0}, step=None)
# mock wandb step
wandb.init().step = 0
# test wandb.init not called if there is a W&B run
wandb.init().log.reset_mock()
wandb.init.reset_mock()
wandb.run = wandb.init()
logger = WandbLogger()
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({'acc': 1.0}, step=3)
# continue training on same W&B run and offset step
wandb.init().step = 3
logger.finalize('success')
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init().log.assert_called_with({'acc': 1.0}, step=6)
# log hyper parameters
logger.log_hyperparams({'test': None, 'nested': {'a': 1}, 'b': [2, 3, 4]})
wandb.init().config.update.assert_called_once_with(
{
'test': 'None',
'nested/a': 1,
'b': [2, 3, 4]
},
allow_val_change=True,
)
# watch a model
logger.watch('model', 'log', 10)
wandb.init().watch.assert_called_once_with('model', log='log', log_freq=10)
# verify warning for logging at a previous step
assert 'Trying to log at a previous step' not in get_warnings(recwarn)
# current step from wandb should be 6 (last logged step)
logger.experiment.step = 6
# logging at step 2 should raise a warning (step_offset is still 3)
logger.log_metrics({'acc': 1.0}, step=2)
assert 'Trying to log at a previous step' in get_warnings(recwarn)
# logging again at step 2 should not display again the same warning
logger.log_metrics({'acc': 1.0}, step=2)
assert 'Trying to log at a previous step' not in get_warnings(recwarn)
assert logger.name == wandb.init().project_name()
assert logger.version == wandb.init().id
def test_wandb_logger_init(wandb):
"""Verify that basic functionality of wandb logger works.
Wandb doesn't work well with pytest so we have to mock it out here.
"""
# test wandb.init called when there is no W&B run
wandb.run = None
logger = WandbLogger(
name="test_name", save_dir="test_save_dir", version="test_id", project="test_project", resume="never"
)
logger.log_metrics({"acc": 1.0})
wandb.init.assert_called_once_with(
name="test_name", dir="test_save_dir", id="test_id", project="test_project", resume="never", anonymous=None
)
wandb.init().log.assert_called_once_with({"acc": 1.0})
# test wandb.init and setting logger experiment externally
wandb.run = None
run = wandb.init()
logger = WandbLogger(experiment=run)
assert logger.experiment
# test wandb.init not called if there is a W&B run
wandb.init().log.reset_mock()
wandb.init.reset_mock()
wandb.run = wandb.init()
logger = WandbLogger()
# verify default resume value
assert logger._wandb_init["resume"] == "allow"
with pytest.warns(UserWarning, match="There is a wandb run already in progress"):
_ = logger.experiment
logger.log_metrics({"acc": 1.0}, step=3)
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({"acc": 1.0, "trainer/global_step": 3})
# continue training on same W&B run and offset step
logger.finalize("success")
logger.log_metrics({"acc": 1.0}, step=6)
wandb.init().log.assert_called_with({"acc": 1.0, "trainer/global_step": 6})
# log hyper parameters
logger.log_hyperparams({"test": None, "nested": {"a": 1}, "b": [2, 3, 4]})
wandb.init().config.update.assert_called_once_with(
{"test": "None", "nested/a": 1, "b": [2, 3, 4]}, allow_val_change=True
)
# watch a model
logger.watch("model", "log", 10, False)
wandb.init().watch.assert_called_once_with("model", log="log", log_freq=10, log_graph=False)
assert logger.name == wandb.init().project_name()
assert logger.version == wandb.init().id
def test_wandb_logger_init(wandb):
"""Verify that basic functionality of wandb logger works.
Wandb doesn't work well with pytest so we have to mock it out here."""
# test wandb.init called when there is no W&B run
wandb.run = None
logger = WandbLogger(
name='test_name', save_dir='test_save_dir', version='test_id', project='test_project', resume='never'
)
logger.log_metrics({'acc': 1.0})
wandb.init.assert_called_once_with(
name='test_name', dir='test_save_dir', id='test_id', project='test_project', resume='never', anonymous=None
)
wandb.init().log.assert_called_once_with({'acc': 1.0})
# test wandb.init and setting logger experiment externally
wandb.run = None
run = wandb.init()
logger = WandbLogger(experiment=run)
assert logger.experiment
# test wandb.init not called if there is a W&B run
wandb.init().log.reset_mock()
wandb.init.reset_mock()
wandb.run = wandb.init()
logger = WandbLogger()
# verify default resume value
assert logger._wandb_init['resume'] == 'allow'
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({'acc': 1.0, 'trainer/global_step': 3})
# continue training on same W&B run and offset step
logger.finalize('success')
logger.log_metrics({'acc': 1.0}, step=6)
wandb.init().log.assert_called_with({'acc': 1.0, 'trainer/global_step': 6})
# log hyper parameters
logger.log_hyperparams({'test': None, 'nested': {'a': 1}, 'b': [2, 3, 4]})
wandb.init().config.update.assert_called_once_with(
{
'test': 'None',
'nested/a': 1,
'b': [2, 3, 4]
},
allow_val_change=True,
)
# watch a model
logger.watch('model', 'log', 10)
wandb.init().watch.assert_called_once_with('model', log='log', log_freq=10)
assert logger.name == wandb.init().project_name()
assert logger.version == wandb.init().id
def test_wandb_logger_init(wandb, recwarn):
"""Verify that basic functionality of wandb logger works.
Wandb doesn't work well with pytest so we have to mock it out here."""
# test wandb.init called when there is no W&B run
wandb.run = None
logger = WandbLogger()
logger.log_metrics({'acc': 1.0})
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({'acc': 1.0})
# test wandb.init not called if there is a W&B run
wandb.init().log.reset_mock()
wandb.init.reset_mock()
wandb.run = wandb.init()
logger = WandbLogger()
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({
'acc': 1.0,
'trainer/global_step': 3
})
# continue training on same W&B run and offset step
logger.finalize('success')
logger.log_metrics({'acc': 1.0}, step=6)
wandb.init().log.assert_called_with({'acc': 1.0, 'trainer/global_step': 6})
# log hyper parameters
logger.log_hyperparams({'test': None, 'nested': {'a': 1}, 'b': [2, 3, 4]})
wandb.init().config.update.assert_called_once_with(
{
'test': 'None',
'nested/a': 1,
'b': [2, 3, 4]
},
allow_val_change=True,
)
# watch a model
logger.watch('model', 'log', 10)
wandb.init().watch.assert_called_once_with('model', log='log', log_freq=10)
assert logger.name == wandb.init().project_name()
assert logger.version == wandb.init().id
def test_wandb_logger(wandb):
"""Verify that basic functionality of wandb logger works.
Wandb doesn't work well with pytest so we have to mock it out here."""
tutils.reset_seed()
logger = WandbLogger(anonymous=True, offline=True)
logger.log_metrics({'acc': 1.0})
wandb.init().log.assert_called_once_with({'acc': 1.0})
wandb.init().log.reset_mock()
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init().log.assert_called_once_with({'global_step': 3, 'acc': 1.0})
logger.log_hyperparams({'test': None})
wandb.init().config.update.assert_called_once_with({'test': None})
logger.watch('model', 'log', 10)
wandb.watch.assert_called_once_with('model', log='log', log_freq=10)
logger.finalize('fail')
wandb.join.assert_called_once_with(1)
wandb.join.reset_mock()
logger.finalize('success')
wandb.join.assert_called_once_with(0)
wandb.join.reset_mock()
wandb.join.side_effect = TypeError
with pytest.raises(TypeError):
logger.finalize('any')
wandb.join.assert_called()
assert logger.name == wandb.init().project_name()
assert logger.version == wandb.init().id
checkpoint = torch.load(path_to_pretrained, map_location="cpu")
model.load_state_dict(checkpoint)
# Additional trainer params
accumulate_grad_batches = args.gradient_acc_batches if args.gradient_acc_batches is not None else None
print("Training", model_name)
print("Initializing the trainer")
exp = wandb.init(project="knowledgeinjection", name=model_name)
wandb.watch(model, log="all")
logger = WandbLogger(name=model_name,
project="knowledgeinjection",
experiment=exp)
trainer = pl.Trainer(callbacks=callbacks,
gpus=args.gpus if args.use_gpu else None,
auto_select_gpus=args.use_gpu,
max_epochs=epochs,
val_check_interval=0.25,
logger=logger,
precision=fp16,
log_every_n_steps=10)
print("Fitting...")
trainer.fit(model, datamodule=dataset)
print("Testing...")
trainer.test(datamodule=dataset)
print("Done!")
logger.finalize("success")
# os.makedirs("models/"+model_name+"/",exist_ok=True)
# trainer.save_checkpoint("models/"+model_name+"/"+model_name+"_trained")
# load data, model
data = HuggingFaceDataModule(config["data"])
model = HuggingFaceModel(config["name"], config["model"], data.tokenizer)
# load experimental setup stuff
callbacks = parse_callbacks(config.get("callbacks", {}))
if not args.fast_dev_run and not args.dry_run:
experiment_logger = WandbLogger(**config.get("logger", {}))
experiment_logger.log_hyperparams(config["model"])
else:
experiment_logger = None
trainer = Trainer.from_argparse_args(args,
logger=experiment_logger,
callbacks=callbacks)
# run experiment
if not args.test:
data.setup('train')
trainer.fit(model, data.train_dataloader(), data.val_dataloader())
else:
data.setup('test')
logger.warning(
"To protect you from yourself, this eval loop loads the validation dataset. Change this manually if you want to actually run on test."
)
trainer.test(test_dataloaders=data.val_dataloader(),
ckpt_path=args.ckpt_path)
if experiment_logger:
experiment_logger.finalize('success')
