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})
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, '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(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
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 main(conf):
# set seed
pl.seed_everything(conf.seed)
# load datamodule
data = HotpotDataModule(conf=conf)
# load model
model_name = conf.model.name + ('_dataset' if conf.model.dataset else '')
if conf.training.train:
if conf.training.from_checkpoint:
model = models[model_name].load_from_checkpoint(
checkpoint_path=os.path.join(
os.path.split(hydra.utils.get_original_cwd())[0],
'outputs', conf.training.from_checkpoint))
else:
model = models[model_name](conf=conf)
else:
model = models[model_name].load_from_checkpoint(
checkpoint_path=os.path.join(
os.path.split(hydra.utils.get_original_cwd())[0], 'outputs',
conf.testing.model_path))
# TRAINER
callbacks = []
# checkpoint callback
checkpoint_callback = ModelCheckpoint(
dirpath=conf.training.model_checkpoint.dirpath,
filename=conf.training.model_checkpoint.filename,
monitor=conf.training.model_checkpoint.monitor,
save_last=conf.training.model_checkpoint.save_last,
save_top_k=conf.training.model_checkpoint.save_top_k)
callbacks.append(checkpoint_callback)
# early stop callback
if conf.training.early_stopping.early_stop:
early_stop_callback = EarlyStopping(
monitor=conf.training.early_stopping.monitor,
patience=conf.training.early_stopping.patience,
mode=conf.training.early_stopping.mode,
)
callbacks.append(early_stop_callback)
# logger
wandb_logger = WandbLogger(name=model_name,
project='neural-question-generation')
# trainer
trainer = pl.Trainer(
accumulate_grad_batches=conf.training.grad_cum,
callbacks=callbacks,
default_root_dir='.',
deterministic=True,
fast_dev_run=conf.debug,
flush_logs_every_n_steps=10,
gpus=(1 if torch.cuda.is_available() else 0),
logger=wandb_logger,
log_every_n_steps=100,
max_epochs=conf.training.max_epochs,
num_sanity_val_steps=0,
reload_dataloaders_every_epoch=True,
# val_check_interval=0.05,
)
# TODO: tune
# train
if conf.training.train:
trainer.fit(model=model, datamodule=data)
# test
if conf.testing.test:
trainer.test(model=model, datamodule=data)
if model_name != 'bert_clf' and model_name != 'bert_sum' and model_name != 'bert_clf+bart_dataset':
results = evaluation_metrics(conf)
wandb_logger.log_metrics(results)
def train(hparams):
EMBEDDING_DIM = 128
USE_AMP = None
NUM_GPUS = hparams.num_gpus
MAX_EPOCHS = 1000
batch_order = 11
dataset = load_node_dataset(hparams.dataset,
hparams.method,
hparams=hparams,
train_ratio=hparams.train_ratio)
METRICS = [
"precision",
"recall",
"f1",
"accuracy",
"top_k" if dataset.multilabel else "ogbn-mag",
]
if hparams.method == "HAN":
USE_AMP = False
model_hparams = {
"embedding_dim": EMBEDDING_DIM,
"batch_size": 2**batch_order,
"num_layers": 2,
"collate_fn": "HAN_batch",
"train_ratio": dataset.train_ratio,
"loss_type": "BINARY_CROSS_ENTROPY"
if dataset.multilabel else "SOFTMAX_CROSS_ENTROPY",
"n_classes": dataset.n_classes,
"lr": 0.001,
}
model = HAN(Namespace(**model_hparams),
dataset=dataset,
metrics=METRICS)
elif hparams.method == "GTN":
USE_AMP = False
model_hparams = {
"embedding_dim": EMBEDDING_DIM,
"num_channels": len(dataset.metapaths),
"num_layers": 2,
"batch_size": 2**batch_order,
"collate_fn": "HAN_batch",
"train_ratio": dataset.train_ratio,
"loss_type": "BINARY_CROSS_ENTROPY"
if dataset.multilabel else "SOFTMAX_CROSS_ENTROPY",
"n_classes": dataset.n_classes,
"lr": 0.001,
}
model = GTN(Namespace(**model_hparams),
dataset=dataset,
metrics=METRICS)
elif hparams.method == "MetaPath2Vec":
USE_AMP = False
model_hparams = {
"embedding_dim": EMBEDDING_DIM,
"walk_length": 50,
"context_size": 7,
"walks_per_node": 5,
"num_negative_samples": 5,
"sparse": True,
"batch_size": 400,
"train_ratio": dataset.train_ratio,
"n_classes": dataset.n_classes,
"lr": 0.01,
}
model = MetaPath2Vec(Namespace(**model_hparams),
dataset=dataset,
metrics=METRICS)
elif hparams.method == "HGT":
USE_AMP = False
model_hparams = {
"embedding_dim": EMBEDDING_DIM,
"num_channels": len(dataset.metapaths),
"n_layers": 2,
"attn_heads": 8,
"attn_dropout": 0.2,
"prev_norm": True,
"last_norm": True,
"nb_cls_dense_size": 0,
"nb_cls_dropout": 0.0,
"use_class_weights": False,
"batch_size": 2**batch_order,
"n_epoch": MAX_EPOCHS,
"train_ratio": dataset.train_ratio,
"loss_type":
"BCE" if dataset.multilabel else "SOFTMAX_CROSS_ENTROPY",
"n_classes": dataset.n_classes,
"collate_fn": "collate_HGT_batch",
"lr": 0.001, # Not used here, defaults to 1e-3
}
model = HGT(Namespace(**model_hparams), dataset, metrics=METRICS)
elif "LATTE" in hparams.method:
USE_AMP = False
num_gpus = 1
if "-1" in hparams.method:
n_layers = 1
elif "-2" in hparams.method:
n_layers = 2
elif "-3" in hparams.method:
n_layers = 3
else:
n_layers = 2
model_hparams = {
"embedding_dim": EMBEDDING_DIM,
"layer_pooling": "concat",
"n_layers": n_layers,
"batch_size": 2**batch_order,
"nb_cls_dense_size": 0,
"nb_cls_dropout": 0.4,
"activation": "relu",
"dropout": 0.2,
"attn_heads": 2,
"attn_activation": "sharpening",
"batchnorm": False,
"layernorm": False,
"edge_sampling": False,
"edge_threshold": 0.5,
"attn_dropout": 0.2,
"loss_type":
"BCE" if dataset.multilabel else "SOFTMAX_CROSS_ENTROPY",
"use_proximity": True if "proximity" in hparams.method else False,
"neg_sampling_ratio": 2.0,
"n_classes": dataset.n_classes,
"use_class_weights": False,
"lr": 0.001,
"momentum": 0.9,
"weight_decay": 1e-2,
}
model_hparams.update(hparams.__dict__)
metrics = [
"precision", "recall", "micro_f1", "macro_f1",
"accuracy" if dataset.multilabel else "ogbn-mag", "top_k"
]
model = LATTENodeClf(Namespace(**model_hparams),
dataset,
collate_fn="neighbor_sampler",
metrics=metrics)
MAX_EPOCHS = 250
wandb_logger = WandbLogger(name=model.name(),
tags=[dataset.name()],
anonymous=True,
project="anon-demo")
wandb_logger.log_hyperparams(model_hparams)
trainer = Trainer(gpus=NUM_GPUS,
distributed_backend='dp' if NUM_GPUS > 1 else None,
max_epochs=MAX_EPOCHS,
stochastic_weight_avg=True,
callbacks=[
EarlyStopping(monitor='val_loss',
patience=10,
min_delta=0.0001,
strict=False)
],
logger=wandb_logger,
weights_summary='top',
amp_level='O1' if USE_AMP and NUM_GPUS > 0 else None,
precision=16 if USE_AMP else 32)
trainer.fit(model)
model.register_hooks()
trainer.test(model)
wandb_logger.log_metrics(
model.clustering_metrics(n_runs=10, compare_node_types=False))
]
if fresh_ckpt_files == []:
time.sleep(0.5)
continue
oldest_ckpt_file = fresh_ckpt_files[0]
print(f'Evaluating checkpoint: {oldest_ckpt_file}')
processing_ckpt_file_name = oldest_ckpt_file.replace('SAVED', 'PROCESSING')
os.rename(os.path.join(save_dir, oldest_ckpt_file),
os.path.join(save_dir, processing_ckpt_file_name))
valid_dict = valid_fn(os.path.join(save_dir, processing_ckpt_file_name),
device=f'cuda:{args.gpus[0]}',
**valid_fn_kwargs)
valid_score = valid_dict[valid_metric]
ckpt_dict = torch.load(os.path.join(save_dir, processing_ckpt_file_name),
map_location='cpu')
valid_dict['global_step'] = ckpt_dict['global_step']
logger.log_metrics(valid_dict)
validated_ckpt_file_name = processing_ckpt_file_name.replace(
'PROCESSING', f'VALID_{float(valid_score):.4f}')
os.rename(os.path.join(save_dir, processing_ckpt_file_name),
os.path.join(save_dir, validated_ckpt_file_name))
print(
f'Evaluation complete, renaming checkpoint to: {validated_ckpt_file_name}'
)
print()
print("getting test set predictions")
with torch.no_grad():
adj_hat_test, y_pred_test, _, _ = model(test_tup[0], test_tup[1])
# print("adj type: {}".format(test_tup[1].flatten().numpy()))
# print("adj_hat type: {}".format(adj_hat_test.flatten().detach().numpy()))
ap_test = average_precision_score(test_tup[1].flatten().numpy(),
adj_hat_test.flatten().detach().numpy())
pred_test = nn.MSELoss()(y_pred_test.flatten(), test_tup[-1])
print("logging test set metrics")
wandb_logger.log_metrics({
'test_energy_mse': pred_test.numpy(),
'test_AP': ap_test
})
print("gathering eval subsets")
eval_tup_list = [
eval_metrics.compute_subsample([test_embed, test_tup[-1]], 10000)[0]
for i in range(8)
]
# trainer.test()
print("getting smoothness vals")
embed_eval_array = np.expand_dims(
np.array([x[0].numpy() for x in eval_tup_list]), 0)
energy_eval_array = np.array([x[1].numpy() for x in eval_tup_list])
print('embed_eval_array shape: {}'.format(embed_eval_array.shape))
print('energy_eval_array shape: {}'.format(energy_eval_array.shape))
