def main(config):
"""
Main function for training LSTMs.
After training, results on validation & test sets are recorded in the specified log_path.
"""
dataset, train_loader, subgraph_loader = get_data(config)
# define logger
Path(config['log_path']).mkdir(parents=True, exist_ok=True)
logger = loggers.TensorBoardLogger(config['log_path'],
version=config['version'])
logger.log_hyperparams(params=config)
# define model
model = Model(config, dataset, train_loader, subgraph_loader)
chkpt = None if config['load'] is None else get_checkpoint_path(
config['load'])
trainer = pl.Trainer(gpus=config['gpus'],
logger=logger,
max_epochs=config['epochs'],
distributed_backend='dp',
precision=16 if config['use_amp'] else 32,
default_root_dir=config['log_path'],
deterministic=True,
resume_from_checkpoint=chkpt,
auto_lr_find=config['auto_lr'],
auto_scale_batch_size=config['auto_bsz'])
trainer.fit(model)
for phase in ['test', 'valid']:
if phase == 'valid':
trainer.eval_split = 'val'
trainer.eval_mask = dataset.data.val_mask
print(phase, trainer.eval_split)
ret = trainer.test()
if isinstance(ret, list):
ret = ret[0]
per_node = ret.pop('per_node')
test_results = ret
res_dir = Path(config['log_path']) / 'default'
if config['version'] is not None:
res_dir = res_dir / config['version']
else:
res_dir = res_dir / ('results_' + str(config['seed']))
print(phase, ':', test_results)
Path(res_dir).mkdir(parents=True, exist_ok=True)
write_json(test_results,
res_dir / f'{phase}_results.json',
sort_keys=True,
verbose=True)
write_pkl(per_node, res_dir / f'{phase}_per_node.pkl')
path_results = Path(config['log_path']) / f'all_{phase}_results.csv'
record_results(path_results, config, test_results)
def main_test(hparams, path_results=None):
"""
main function to load and evaluate a trained model.
"""
assert (hparams['load'] is not None) and (hparams['phase'] is not None)
phase = hparams['phase']
log_dir = hparams['load']
# Load trained model
print(f'Loading from {log_dir} to evaluate {phase} data.')
model, config, dataset, train_loader, subgraph_loader = Model.load_model(
log_dir,
multi_gpu=hparams['multi_gpu'],
num_workers=hparams['num_workers'])
trainer = pl.Trainer(gpus=hparams['gpus'],
logger=None,
max_epochs=hparams['epochs'],
default_root_dir=hparams['log_path'],
deterministic=True)
# Evaluate the model
if phase == 'valid':
trainer.eval_split = 'val'
trainer.eval_mask = dataset.data.val_mask
print(phase, trainer.eval_split)
test_results = trainer.test(model)
if isinstance(test_results, list):
test_results = test_results[0]
per_node = test_results.pop('per_node')
print(phase, ':', test_results)
# Save evaluation results
results_path = Path(log_dir) / f'{phase}_results.json'
write_json(test_results, results_path, sort_keys=True, verbose=True)
write_pkl(per_node, Path(log_dir) / f'{phase}_per_node.pkl')
if path_results is None:
path_results = Path(log_dir).parent / 'results.csv'
tmp = {'version': hparams['version']}
tmp = {**tmp, **config}
record_results(path_results, tmp, test_results)
def main_test(hparams, path_results=None):
"""
main function to load and evaluate a trained model.
"""
assert (hparams['load'] is not None) and (hparams['phase'] is not None)
phase = hparams['phase']
log_dir = hparams['load']
# Load trained model
print(f'Loading from {log_dir} to evaluate {phase} data.')
model, config, loaderDict, collate = DynamicGraphModel.load_model(log_dir, \
data_dir=hparams['data_dir'],
multi_gpu=hparams['multi_gpu'], num_workers=hparams['num_workers'])
trainer = pl.Trainer(
gpus=hparams['gpus'],
logger=None,
max_epochs=hparams['epochs'],
default_root_dir=hparams['log_path'],
deterministic=True
)
# Evaluate the model
test_dataloader = DataLoader(loaderDict[phase], collate_fn=collate, batch_size=config['batch_size'], num_workers=config['num_workers'], shuffle=False)
test_results = trainer.test(model, test_dataloaders=test_dataloader)
if isinstance(test_results, list):
test_results = test_results[0]
per_node = test_results.pop('per_node')
print(phase, ':', test_results)
# Save evaluation results
results_path = Path(log_dir) / f'{phase}_results.json'
write_json(test_results, results_path, sort_keys=True, verbose=True)
write_pkl(per_node, Path(log_dir) / f'{phase}_per_node.pkl')
if path_results is None:
path_results = Path(log_dir).parent / 'results.csv'
tmp = {'version': hparams['version']}
tmp = {**tmp, **config}
record_results(path_results, tmp, test_results)