def main():
tvt_pairs_dict = load_pair_tvt_splits()
orig_dataset = load_dataset(FLAGS.dataset, 'all', FLAGS.node_feats,
FLAGS.edge_feats)
orig_dataset, num_node_feat = encode_node_features(dataset=orig_dataset)
num_interaction_edge_feat = encode_edge_features(
orig_dataset.interaction_combo_nxgraph, FLAGS.hyper_eatts)
for i, (train_pairs, val_pairs, test_pairs) in \
enumerate(zip(tvt_pairs_dict['train'],
tvt_pairs_dict['val'],
tvt_pairs_dict['test'])):
fold_num = i + 1
if FLAGS.cross_val and FLAGS.run_only_on_fold != -1 and FLAGS.run_only_on_fold != fold_num:
continue
set_seed(FLAGS.random_seed + 5)
print(f'======== FOLD {fold_num} ========')
saver = Saver(fold=fold_num)
dataset = deepcopy(orig_dataset)
train_data, val_data, test_data, val_pairs, test_pairs, _ = \
load_pairs_to_dataset(num_node_feat, num_interaction_edge_feat,
train_pairs, val_pairs, test_pairs,
dataset)
print('========= Training... ========')
if FLAGS.load_model is not None:
print('loading models: {}'.format(FLAGS.load_model))
trained_model = Model(train_data).to(FLAGS.device)
trained_model.load_state_dict(torch.load(
FLAGS.load_model, map_location=FLAGS.device),
strict=False)
print('models loaded')
print(trained_model)
else:
train(train_data, val_data, val_pairs, saver, fold_num=fold_num)
trained_model = saver.load_trained_model(train_data)
if FLAGS.save_model:
saver.save_trained_model(trained_model)
print('======== Testing... ========')
if FLAGS.lower_level_layers and FLAGS.higher_level_layers:
_get_initial_embd(test_data, trained_model)
test_data.dataset.init_interaction_graph_embds(device=FLAGS.device)
elif FLAGS.higher_level_layers and not FLAGS.lower_level_layers:
test_data.dataset.init_interaction_graph_embds(device=FLAGS.device)
if FLAGS.save_final_node_embeddings and 'gmn' not in FLAGS.model:
with torch.no_grad():
trained_model = trained_model.to(FLAGS.device)
trained_model.eval()
if FLAGS.higher_level_layers:
batch_data = model_forward(trained_model,
test_data,
is_train=False)
trained_model.use_layers = "higher_no_eval_layers"
outs = trained_model(batch_data)
else:
outs = _get_initial_embd(test_data, trained_model)
trained_model.use_layers = 'all'
saver.save_graph_embeddings_mat(outs.cpu().detach().numpy(),
test_data.dataset.id_map,
test_data.dataset.gs_map)
if FLAGS.higher_level_layers:
batch_data.restore_interaction_nxgraph()
test(trained_model, test_data, test_pairs, saver, fold_num)
overall_time = convert_long_time_to_str(time() - t)
print(overall_time)
print(saver.get_log_dir())
print(basename(saver.get_log_dir()))
saver.save_overall_time(overall_time)
saver.close()
if FLAGS.cross_val and COMET_EXPERIMENT:
results = aggregate_comet_results_from_folds(
COMET_EXPERIMENT, FLAGS.num_folds, FLAGS.dataset,
FLAGS.eval_performance_by_degree)
COMET_EXPERIMENT.log_metrics(results, prefix='aggr')
args.num_channels = datasets['train'][0][0].shape[0]
try:
args.num_classes = datasets['train'].num_classes
except:
try:
args.num_classes = len(datasets['train'].classes)
except:
pass
assert args.mode == 'auto_encoder' or args.num_classes is not None, "couldn't get number of classes"
# Define saver
saver = Saver(args.logdir,
args,
sub_dirs=list(datasets.keys()),
tag=args.tag)
# Add saver to args (e.g. visualizing weights, outputs)
args.saver = saver
# Save splits
# TODO adapt this to non-CIFAR dataset
saver.save_data({'train': train_idxs, 'val': val_idxs}, 'split_idxs')
# Define trainer
trainer_module = getattr(
trainers, "ae_trainer" if args.mode == 'auto_encoder' else 'trainer')
trainer_class = getattr(trainer_module, 'Trainer')
trainer = trainer_class(args)
# Run training
model, metrics = trainer.train(datasets)
# Close saver
saver.close()
