tgan, val_src_l,
val_dst_l, val_ts_l,
val_label_l)
epoch_bar.update()
epoch_bar.set_postfix(acc=val_acc, f1=val_f1, auc=val_auc)
if early_stopper.early_stop_check(val_auc):
break
else:
torch.save(tgan.state_dict(), get_checkpoint_path(epoch))
logger.info('No improvment over {} epochs, stop training'.format(
early_stopper.max_round))
logger.info(f'Loading the best model at epoch {early_stopper.best_epoch}')
best_model_path = get_checkpoint_path(early_stopper.best_epoch)
tgan.load_state_dict(torch.load(best_model_path))
logger.info(
f'Loaded the best model at epoch {early_stopper.best_epoch} for inference')
tgan.eval()
# testing phase use all information
tgan.ngh_finder = full_ngh_finder
_, _, _, val_auc = eval_one_epoch('val for old nodes', tgan, val_src_l,
val_dst_l, val_ts_l, val_label_l)
test_acc, test_ap, test_f1, test_auc = eval_one_epoch('test for old nodes',
tgan, test_src_l,
test_dst_l, test_ts_l,
test_label_l)
logger.info('Test statistics: acc: {:.4f}, f1:{:.4f} auc: {:.4f}'.format(
test_acc, test_f1, test_auc))
# if early_stopper.early_stop_check(val_mrr):
# logger.info('No improvment over {} epochs, stop training'.format(early_stopper.max_round))
# logger.info(f'Loading the best model at epoch {early_stopper.best_epoch}')
# best_model_path = get_checkpoint_path(early_stopper.best_epoch)
# tgan.load_state_dict(torch.load(best_model_path))
# logger.info(f'Loaded the best model at epoch {early_stopper.best_epoch} for inference')
# tgan.eval()
# break
# else:
torch.save(tgan.state_dict(), get_checkpoint_path(epoch))
# end = time.time()
# print(f'total consume {end - start}')
torch.save(tgan.state_dict(), MODEL_SAVE_PATH)
tgan.load_state_dict(torch.load(MODEL_SAVE_PATH))
# testing phase use all information
tgan.ngh_finder = full_ngh_finder
test_mrr, test_rec10 = eval_one_epoch('test for old nodes', tgan, test_rand_sampler, test_src_l,
test_dst_l, test_ts_l, test_label_l)
nn_test_mrr, nn_test_rec10 = eval_one_epoch('test for new nodes', tgan, nn_test_rand_sampler, nn_test_src_l,
nn_test_dst_l, nn_test_ts_l, nn_test_label_l)
watcher.insert_batch_results([test_mrr, test_rec10, nn_test_mrr, nn_test_rec10])
logger.info('Test statistics: Old nodes -- acc: {}, auc: {}'.format(test_mrr, test_rec10))
logger.info('Test statistics: New nodes -- acc: {}, auc: {}'.format(nn_test_mrr, nn_test_rec10))
logger.info('Saving TGAN model')
logger.info('TGAN models saved')