def main_fixed_mask(args):
device = torch.device("cuda:" + str(args.device))
dataset = PygNodePropPredDataset(name=args.dataset)
data = dataset[0]
split_idx = dataset.get_idx_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
y_true = data.y.to(device)
train_idx = split_idx['train'].to(device)
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
if args.self_loop:
edge_index = add_self_loops(edge_index, num_nodes=data.num_nodes)[0]
args.in_channels = data.x.size(-1)
args.num_tasks = dataset.num_classes
model = DeeperGCN(args).to(device)
pruning.add_mask(model, args.num_layers)
for name, param in model.named_parameters():
if 'mask' in name:
param.requires_grad = False
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {'highest_valid': 0, 'final_train': 0, 'final_test': 0, 'highest_train': 0, 'epoch': 0}
start_epoch = 1
for epoch in range(start_epoch, args.epochs + 1):
epoch_loss = train_fixed(model, x, edge_index, y_true, train_idx, optimizer, args)
result = test(model, x, edge_index, y_true, split_idx, evaluator)
train_accuracy, valid_accuracy, test_accuracy = result
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
results['epoch'] = epoch
print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'Baseline (FIX Mask) Epoch:[{}/{}]\t LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}]'
.format(epoch, args.epochs, epoch_loss, train_accuracy * 100,
valid_accuracy * 100,
test_accuracy * 100,
results['final_test'] * 100,
results['epoch']))
print("=" * 120)
print("syd final: Baseline, Train:[{:.2f}] Best Val:[{:.2f}] at epoch:[{}] | Final Test Acc:[{:.2f}]"
.format( results['final_train'] * 100,
results['highest_valid'] * 100,
results['epoch'],
results['final_test'] * 100))
print("=" * 120)
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygNodePropPredDataset(name=args.dataset, root=args.data_folder)
graph = dataset[0]
adj = SparseTensor(row=graph.edge_index[0], col=graph.edge_index[1])
if args.self_loop:
adj = adj.set_diag()
graph.edge_index = add_self_loops(edge_index=graph.edge_index,
num_nodes=graph.num_nodes)[0]
split_idx = dataset.get_idx_split()
train_idx = split_idx["train"].tolist()
evaluator = Evaluator(args.dataset)
sub_dir = 'random-train_{}-full_batch_test'.format(args.cluster_number)
logging.info(sub_dir)
log_dir = os.path.join(args.save, "tensorboard/")
writer = SummaryWriter(log_dir=log_dir)
args.in_channels = graph.x.size(-1)
args.num_tasks = dataset.num_classes
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
# generate batches
parts = random_partition_graph(graph.num_nodes,
cluster_number=args.cluster_number)
data = generate_sub_graphs(adj,
parts,
cluster_number=args.cluster_number)
# epoch_loss = train(data, model, graph.x, graph.y, train_idx, optimizer, device)
epoch_loss = train_flag(data, model, graph.x, graph.y, train_idx,
optimizer, device, args)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
writer.add_scalar("Loss/train", epoch_loss, epoch)
model.print_params(epoch=epoch)
# if epoch % args.eval_epochs == 0:
# save_ckpt(model, optimizer,
# round(epoch_loss, 4), epoch,
# args.model_save_path,
# sub_dir, name_post=f'epoch{epoch}')
# logging.info(f"Epoch {epoch}, saved model to checkpoint folder {args.model_save_path}")
if epoch % args.eval_epochs == 0:
# save_ckpt(model, optimizer,
# round(epoch_loss, 4), epoch,
# args.model_save_path,
# sub_dir, name_post='final')
# logging.info(f"Saved model to checkpoint folder {args.model_save_path}")
logging.info(f'---- Evaluating at epoch {epoch} ----')
res = test_with_partition(model,
graph,
adj,
split_idx,
num_clusters=args.eval_cluster_number,
partition_method=args.partition_method,
evaluator=evaluator,
device=device)
# result = test(model, graph.x, graph.edge_index, graph.y, split_idx, evaluator)
logging.info(res)
logging.info(f"---------------------------------")
train_accuracy, valid_accuracy, test_accuracy = res[
"train_acc"], res["valid_acc"], res["test_acc"]
writer.add_scalar("Acc/train", train_accuracy)
writer.add_scalar("Acc/dev", valid_accuracy)
if train_accuracy > results['highest_train']:
results['highest_train'] = train_accuracy
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
logging.info(
f"Saved better model to checkpoint folder {args.model_save_path}"
)
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
dic = torch.load('/Data/leo/Pixel2Mesh_3d/2dTo3d/wandb/run-20210126_131356-1bkxbzmf/files/dic_checkpoint')
gcn1.load_state_dict(dic['gcn_state_dict'][0])
gcn1.train()
gcn2.load_state_dict(dic['gcn_state_dict'][1])
gcn2.train()
gcn3.load_state_dict(dic['gcn_state_dict'][2])
gcn3.train()
image_model.load_state_dict(dic['cnn_state_dic'])
image_model.train()
fe.load_state_dict(dic['fe_state_dic'])
fe.train()
#optimizer.load_state_dict(dic['optimizer_state_dict'])
#epoch = checkpoint['epoch']
#loss = checkpoint['loss']
optimizer = torch.optim.Adam(list(gcn1.parameters()) +
list(gcn2.parameters()) +
list(gcn3.parameters()) +
list(image_model.parameters()), lr=5e-5)
#scheduler = MultiStepLR(optimizer, milestones=[5, 10, 30, 80], gamma=0.3)
w_chamfer = 1.0
# Weight for mesh edge loss
w_edge = 0.1#1.0
# Weight for mesh normal consistency
w_normal = 1.6e-4#0.01
# Weight for mesh laplacian smoothing
w_laplacian = 0.3#0.1
# Plot period for the losses
def main_fixed_mask(args, imp_num, final_state_dict=None, resume_train_ckpt=None):
device = torch.device("cuda:" + str(args.device))
dataset = OGBNDataset(dataset_name=args.dataset)
nf_path = dataset.extract_node_features(args.aggr)
args.num_tasks = dataset.num_tasks
args.nf_path = nf_path
evaluator = Evaluator(args.dataset)
criterion = torch.nn.BCEWithLogitsLoss()
valid_data_list = []
for i in range(args.num_evals):
parts = dataset.random_partition_graph(dataset.total_no_of_nodes, cluster_number=args.valid_cluster_number)
valid_data = dataset.generate_sub_graphs(parts, cluster_number=args.valid_cluster_number)
valid_data_list.append(valid_data)
print("-" * 120)
model = DeeperGCN(args).to(device)
pruning.add_mask(model)
if final_state_dict is not None:
pruning.retrain_operation(dataset, model, final_state_dict)
adj_spar, wei_spar = pruning.print_sparsity(dataset, model)
for name, param in model.named_parameters():
if 'mask' in name:
param.requires_grad = False
optimizer = optim.Adam(model.parameters(), lr=args.lr)
results = {'highest_valid': 0, 'final_train': 0, 'final_test': 0, 'highest_train': 0, 'epoch':0}
results['adj_spar'] = adj_spar
results['wei_spar'] = wei_spar
start_epoch = 1
if resume_train_ckpt:
dataset.adj = resume_train_ckpt['adj']
start_epoch = resume_train_ckpt['epoch']
rewind_weight_mask = resume_train_ckpt['rewind_weight_mask']
ori_model_dict = model.state_dict()
over_lap = {k : v for k, v in resume_train_ckpt['model_state_dict'].items() if k in ori_model_dict.keys()}
ori_model_dict.update(over_lap)
model.load_state_dict(ori_model_dict)
print("Resume at IMP:[{}] epoch:[{}] len:[{}/{}]!".format(imp_num, resume_train_ckpt['epoch'], len(over_lap.keys()), len(ori_model_dict.keys())))
optimizer.load_state_dict(resume_train_ckpt['optimizer_state_dict'])
adj_spar, wei_spar = pruning.print_sparsity(dataset, model)
for epoch in range(start_epoch, args.epochs + 1):
# do random partition every epoch
t0 = time.time()
train_parts = dataset.random_partition_graph(dataset.total_no_of_nodes, cluster_number=args.cluster_number)
data = dataset.generate_sub_graphs(train_parts, cluster_number=args.cluster_number, ifmask=True)
epoch_loss = train.train_fixed(data, dataset, model, optimizer, criterion, device, args)
result = train.multi_evaluate(valid_data_list, dataset, model, evaluator, device)
train_result = result['train']['rocauc']
valid_result = result['valid']['rocauc']
test_result = result['test']['rocauc']
if valid_result > results['highest_valid']:
results['highest_valid'] = valid_result
results['final_train'] = train_result
results['final_test'] = test_result
results['epoch'] = epoch
final_state_dict = pruning.save_all(dataset,
model,
None,
optimizer,
imp_num,
epoch,
args.model_save_path,
'IMP{}_fixed_ckpt'.format(imp_num))
epoch_time = (time.time() - t0) / 60
print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'IMP:[{}] (FIX Mask) Epoch[{}/{}] LOSS[{:.4f}] Train[{:.2f}] Valid[{:.2f}] Test[{:.2f}] | Update Test[{:.2f}] at epoch[{}] | Adj[{:.2f}%] Wei[{:.2f}%] Time[{:.2f}min]'
.format(imp_num, epoch, args.epochs, epoch_loss, train_result * 100,
valid_result * 100,
test_result * 100,
results['final_test'] * 100,
results['epoch'],
results['adj_spar'] * 100,
results['wei_spar'] * 100,
epoch_time))
print("=" * 120)
print("INFO final: IMP:[{}], Train:[{:.2f}] Best Val:[{:.2f}] at epoch:[{}] | Final Test Acc:[{:.2f}] | Adj:[{:.2f}%] Wei:[{:.2f}%]"
.format(imp_num, results['final_train'] * 100,
results['highest_valid'] * 100,
results['epoch'],
results['final_test'] * 100,
results['adj_spar'] * 100,
results['wei_spar'] * 100))
print("=" * 120)
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
sub_dir = 'BS_{}'.format(args.batch_size)
if args.not_extract_node_feature:
dataset = PygGraphPropPredDataset(name=args.dataset,
transform=add_zeros)
else:
if args.aggr == 'add':
dataset = PygGraphPropPredDataset(
name=args.dataset, transform=extract_node_feature_add)
elif args.aggr == 'mean':
dataset = PygGraphPropPredDataset(
name=args.dataset, transform=extract_node_feature_mean)
elif args.aggr == 'max':
dataset = PygGraphPropPredDataset(
name=args.dataset, transform=extract_node_feature_max)
else:
raise Exception('Unknown Aggregation Type')
sub_dir = sub_dir + '-NF_{}'.format(args.aggr)
args.num_tasks = dataset.num_classes
evaluator = Evaluator(args.dataset)
logging.info('%s' % args)
split_idx = dataset.get_idx_split()
train_loader = DataLoader(dataset[split_idx["train"]],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
valid_loader = DataLoader(dataset[split_idx["valid"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
test_loader = DataLoader(dataset[split_idx["test"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
criterion = torch.nn.CrossEntropyLoss()
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
evaluate = True
for epoch in range(1, args.epochs + 1):
logging.info("=====Epoch {}".format(epoch))
logging.info('Training...')
epoch_loss = train(model, device, train_loader, optimizer, criterion)
if args.num_layers > args.num_layers_threshold:
if epoch % args.eval_steps != 0:
evaluate = False
else:
evaluate = True
model.print_params(epoch=epoch)
if evaluate:
logging.info('Evaluating...')
train_accuracy = eval(model, device, train_loader, evaluator)
valid_accuracy = eval(model, device, valid_loader, evaluator)
test_accuracy = eval(model, device, test_loader, evaluator)
logging.info({
'Train': train_accuracy,
'Validation': valid_accuracy,
'Test': test_accuracy
})
if train_accuracy > results['highest_train']:
results['highest_train'] = train_accuracy
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
def main_fixed_mask(args,
imp_num,
adj_percent,
wei_percent,
resume_train_ckpt=None):
device = torch.device("cuda:" + str(args.device))
dataset = PygNodePropPredDataset(name=args.dataset)
data = dataset[0]
split_idx = dataset.get_idx_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
y_true = data.y.to(device)
train_idx = split_idx['train'].to(device)
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
if args.self_loop:
edge_index = add_self_loops(edge_index, num_nodes=data.num_nodes)[0]
args.in_channels = data.x.size(-1)
args.num_tasks = dataset.num_classes
model = DeeperGCN(args).to(device)
pruning.add_mask(model, args.num_layers)
pruning.random_pruning(model, args, adj_percent, wei_percent)
adj_spar, wei_spar = pruning.print_sparsity(model, args)
for name, param in model.named_parameters():
if 'mask' in name:
param.requires_grad = False
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0,
'epoch': 0
}
results['adj_spar'] = adj_spar
results['wei_spar'] = wei_spar
start_epoch = 1
if resume_train_ckpt:
start_epoch = resume_train_ckpt['epoch']
ori_model_dict = model.state_dict()
over_lap = {
k: v
for k, v in resume_train_ckpt['model_state_dict'].items()
if k in ori_model_dict.keys()
}
ori_model_dict.update(over_lap)
model.load_state_dict(ori_model_dict)
print("(RP FIXED MASK) Resume at epoch:[{}] len:[{}/{}]!".format(
resume_train_ckpt['epoch'], len(over_lap.keys()),
len(ori_model_dict.keys())))
optimizer.load_state_dict(resume_train_ckpt['optimizer_state_dict'])
adj_spar, wei_spar = pruning.print_sparsity(model, args)
for epoch in range(start_epoch, args.epochs + 1):
epoch_loss = train_fixed(model, x, edge_index, y_true, train_idx,
optimizer, args)
result = test(model, x, edge_index, y_true, split_idx, evaluator)
train_accuracy, valid_accuracy, test_accuracy = result
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
results['epoch'] = epoch
pruning.save_all(model, None, optimizer, imp_num, epoch,
args.model_save_path,
'RP{}_fixed_ckpt'.format(imp_num))
print(
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'RP:[{}] (FIX Mask) Epoch:[{}/{}]\t LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}]'
.format(imp_num, epoch, args.epochs, epoch_loss, train_accuracy *
100, valid_accuracy * 100, test_accuracy *
100, results['final_test'] * 100, results['epoch']))
print("=" * 120)
print(
"syd final: RP:[{}], Train:[{:.2f}] Best Val:[{:.2f}] at epoch:[{}] | Final Test Acc:[{:.2f}] Adj:[{:.2f}%] Wei:[{:.2f}%]"
.format(imp_num, results['final_train'] * 100,
results['highest_valid'] * 100, results['epoch'],
results['final_test'] * 100, results['adj_spar'],
results['wei_spar']))
print("=" * 120)
def main(args):
device = torch.device("cuda:" + str(args.device))
dataset = PygLinkPropPredDataset(name=args.dataset)
data = dataset[0]
# Data(edge_index=[2, 2358104], edge_weight=[2358104, 1], edge_year=[2358104, 1], x=[235868, 128])
split_edge = dataset.get_edge_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
edge_index = data.edge_index.to(device)
args.in_channels = data.x.size(-1)
args.num_tasks = 1
model = DeeperGCN(args).to(device)
predictor = LinkPredictor(args).to(device)
logging.info(model)
logging.info(predictor)
optimizer = torch.optim.Adam(list(model.parameters()) +
list(predictor.parameters()),
lr=args.lr)
results = {}
keys = ['highest_valid', 'final_train', 'final_test', 'highest_train']
hits = ['Hits@10', 'Hits@50', 'Hits@100']
for key in keys:
results[key] = {k: 0 for k in hits}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
epoch_loss = train.train(model, predictor, x, edge_index, split_edge,
optimizer, args.batch_size)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
result = train.test(model, predictor, x, edge_index, split_edge,
evaluator, args.batch_size)
for k in hits:
# return a tuple
train_result, valid_result, test_result = result[k]
if train_result > results['highest_train'][k]:
results['highest_train'][k] = train_result
if valid_result > results['highest_valid'][k]:
results['highest_valid'][k] = valid_result
results['final_train'][k] = train_result
results['final_test'][k] = test_result
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
k,
name_post='valid_best')
save_ckpt(predictor,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
k,
name_post='valid_best_link_predictor')
logging.info(result)
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
time_used = 'Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time)))
logging.info(time_used)
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
sub_dir = 'BS_{}-NF_{}'.format(args.batch_size, args.feature)
dataset = PygGraphPropPredDataset(name=args.dataset)
args.num_tasks = dataset.num_tasks
logging.info('%s' % args)
if args.feature == 'full':
pass
elif args.feature == 'simple':
print('using simple feature')
# only retain the top two node/edge features
dataset.data.x = dataset.data.x[:, :2]
dataset.data.edge_attr = dataset.data.edge_attr[:, :2]
evaluator = Evaluator(args.dataset)
split_idx = dataset.get_idx_split()
train_loader = DataLoader(dataset[split_idx["train"]],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
valid_loader = DataLoader(dataset[split_idx["valid"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
test_loader = DataLoader(dataset[split_idx["test"]],
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
logging.info("=====Epoch {}".format(epoch))
logging.info('Training...')
# epoch_loss = train(model, device, train_loader, optimizer, dataset.task_type)
epoch_loss = train_flag(model, device, train_loader, optimizer,
dataset.task_type, args)
logging.info('Evaluating...')
train_result = eval(model, device, train_loader,
evaluator)[dataset.eval_metric]
valid_result = eval(model, device, valid_loader,
evaluator)[dataset.eval_metric]
test_result = eval(model, device, test_loader,
evaluator)[dataset.eval_metric]
logging.info({
'Train': train_result,
'Validation': valid_result,
'Test': test_result
})
model.print_params(epoch=epoch)
if train_result > results['highest_train']:
results['highest_train'] = train_result
if valid_result > results['highest_valid']:
results['highest_valid'] = valid_result
results['final_train'] = train_result
results['final_test'] = test_result
# save_ckpt(model, optimizer,
# round(epoch_loss, 4), epoch,
# args.model_save_path,
# sub_dir, name_post='valid_best')
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device("cpu")
logging.info('%s' % device)
dataset = OGBNDataset(dataset_name=args.dataset)
# extract initial node features
nf_path = dataset.extract_node_features(args.aggr)
args.num_tasks = dataset.num_tasks
args.nf_path = nf_path
logging.info('%s' % args)
evaluator = Evaluator(args.dataset)
criterion = torch.nn.BCEWithLogitsLoss()
valid_data_list = []
for i in range(args.num_evals):
parts = dataset.random_partition_graph(
dataset.total_no_of_nodes,
cluster_number=args.valid_cluster_number)
valid_data = dataset.generate_sub_graphs(
parts, cluster_number=args.valid_cluster_number)
valid_data_list.append(valid_data)
sub_dir = 'random-train_{}-test_{}-num_evals_{}'.format(
args.cluster_number, args.valid_cluster_number, args.num_evals)
logging.info(sub_dir)
model = DeeperGCN(args).to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
# do random partition every epoch
train_parts = dataset.random_partition_graph(
dataset.total_no_of_nodes, cluster_number=args.cluster_number)
data = dataset.generate_sub_graphs(train_parts,
cluster_number=args.cluster_number,
ifmask=True)
epoch_loss = train(data, dataset, model, optimizer, criterion, device)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
model.print_params(epoch=epoch)
result = multi_evaluate(valid_data_list, dataset, model, evaluator,
device)
if epoch % 5 == 0:
logging.info('%s' % result)
train_result = result['train']['rocauc']
valid_result = result['valid']['rocauc']
test_result = result['test']['rocauc']
if valid_result > results['highest_valid']:
results['highest_valid'] = valid_result
results['final_train'] = train_result
results['final_test'] = test_result
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
if train_result > results['highest_train']:
results['highest_train'] = train_result
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
def main():
args = ArgsInit().save_exp()
if args.use_tensor_board:
writer = SummaryWriter(log_dir=args.save)
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygLinkPropPredDataset(name=args.dataset)
data = dataset[0]
# Data(edge_index=[2, 2358104], edge_weight=[2358104, 1], edge_year=[2358104, 1], x=[235868, 128])
split_edge = dataset.get_edge_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
edge_index = data.edge_index.to(device)
args.in_channels = data.x.size(-1)
args.num_tasks = 1
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
predictor = LinkPredictor(args).to(device)
logging.info(model)
logging.info(predictor)
optimizer = torch.optim.Adam(list(model.parameters()) +
list(predictor.parameters()),
lr=args.lr)
results = {}
keys = ['highest_valid', 'final_train', 'final_test', 'highest_train']
hits = ['Hits@10', 'Hits@50', 'Hits@100']
for key in keys:
results[key] = {k: 0 for k in hits}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
epoch_loss = train(model, predictor, x, edge_index, split_edge,
optimizer, args.batch_size)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
model.print_params(epoch=epoch)
result = test(model, predictor, x, edge_index, split_edge, evaluator,
args.batch_size)
for k in hits:
# return a tuple
train_result, valid_result, test_result = result[k]
if args.use_tensor_board and k == 'Hits@50':
writer.add_scalar('stats/train_loss', epoch_loss, epoch)
writer.add_scalar('stats/train_Hits@50', train_result, epoch)
writer.add_scalar('stats/valid_Hits@50', valid_result, epoch)
writer.add_scalar('stats/test_Hits@50', test_result, epoch)
if train_result > results['highest_train'][k]:
results['highest_train'][k] = train_result
if valid_result > results['highest_valid'][k]:
results['highest_valid'][k] = valid_result
results['final_train'][k] = train_result
results['final_test'][k] = test_result
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
k,
name_post='valid_best')
save_ckpt(predictor,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
k,
name_post='valid_best_link_predictor')
logging.info(result)
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
time_used = 'Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time)))
logging.info(time_used)
#plot_pointcloud(trg_mesh, "Target mesh") #plot_pointcloud(src_mesh, "Source mesh") plt.show() # We will learn to deform the source mesh by offsetting its vertices # The shape of the deform parameters is equal to the total number of vertices in src_mesh# #deform_verts = torch.full(block1.verts_packed().shape, 0.0, device=device, requires_grad=True) #model = GCN() #print(device) model = DeeperGCN(3, 128, 3, 5) model.to(device) # The optimizer optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9) #optimizer = torch.optim.SGD([deform_verts], lr=0.2, momentum=0.9) # Number of optimization steps Niter = 3001 # Weight for the chamfer loss w_chamfer = 1.0 # Weight for mesh edge loss w_edge = 0.1 #1.0 # Weight for mesh normal consistency w_normal = 1.6e-4 #0.01 # Weight for mesh laplacian smoothing w_laplacian = 0.3 #0.1 # Plot period for the losses plot_period = 250
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygNodePropPredDataset(name=args.dataset)
graph = dataset[0]
adj = SparseTensor(row=graph.edge_index[0], col=graph.edge_index[1])
if args.self_loop:
adj = adj.set_diag()
graph.edge_index = add_self_loops(edge_index=graph.edge_index,
num_nodes=graph.num_nodes)[0]
split_idx = dataset.get_idx_split()
train_idx = split_idx["train"].tolist()
evaluator = Evaluator(args.dataset)
sub_dir = 'random-train_{}-full_batch_test'.format(args.cluster_number)
logging.info(sub_dir)
args.in_channels = graph.x.size(-1)
args.num_tasks = dataset.num_classes
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
# generate batches
parts = random_partition_graph(graph.num_nodes,
cluster_number=args.cluster_number)
data = generate_sub_graphs(adj,
parts,
cluster_number=args.cluster_number)
epoch_loss = train(data, model, graph.x, graph.y, train_idx, optimizer,
device)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
model.print_params(epoch=epoch)
if epoch == args.epochs:
result = test(model, graph.x, graph.edge_index, graph.y, split_idx,
evaluator)
logging.info(result)
train_accuracy, valid_accuracy, test_accuracy = result
if train_accuracy > results['highest_train']:
results['highest_train'] = train_accuracy
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
def main_get_mask(args, imp_num):
device = torch.device("cuda:" + str(args.device))
dataset = PygLinkPropPredDataset(name=args.dataset)
data = dataset[0]
# Data(edge_index=[2, 2358104], edge_weight=[2358104, 1], edge_year=[2358104, 1], x=[235868, 128])
split_edge = dataset.get_edge_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
edge_index = data.edge_index.to(device)
args.in_channels = data.x.size(-1)
args.num_tasks = 1
model = DeeperGCN(args).to(device)
pruning.add_mask(model, args)
for name, param in model.named_parameters():
if 'mask' in name:
param.requires_grad = False
predictor = LinkPredictor(args).to(device)
optimizer = torch.optim.Adam(list(model.parameters()) +
list(predictor.parameters()),
lr=args.lr)
results = {'epoch': 0}
keys = ['highest_valid', 'final_train', 'final_test', 'highest_train']
hits = ['Hits@10', 'Hits@50', 'Hits@100']
for key in keys:
results[key] = {k: 0 for k in hits}
start_epoch = 1
for epoch in range(start_epoch, args.mask_epochs + 1):
t0 = time.time()
epoch_loss = train.train_fixed(model, predictor, x, edge_index,
split_edge, optimizer, args.batch_size,
args)
result = train.test(model, predictor, x, edge_index, split_edge,
evaluator, args.batch_size, args)
k = 'Hits@50'
train_result, valid_result, test_result = result[k]
if train_result > results['highest_train'][k]:
results['highest_train'][k] = train_result
if valid_result > results['highest_valid'][k]:
results['highest_valid'][k] = valid_result
results['final_train'][k] = train_result
results['final_test'][k] = test_result
results['epoch'] = epoch
epoch_time = (time.time() - t0) / 60
print(
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'IMP:[{}] (GET Mask) Epoch:[{}/{}] LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}] Time:[{:.2f}min]'
.format(imp_num, epoch, args.mask_epochs, epoch_loss,
train_result * 100, valid_result * 100, test_result *
100, results['final_test'][k] *
100, results['epoch'], epoch_time))
print('-' * 100)
print(
"syd : IMP:[{}] (FIX Mask) Final Result Train:[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}]"
.format(imp_num, results['final_train'][k] * 100,
results['highest_valid'][k] * 100,
results['final_test'][k] * 100))
print('-' * 100)
def main():
EPOCHS = 1
NUMBER_OF_SUBGRAPHS = 10
CKPT_PATH = f'neighbor_deeper_num_{NUMBER_OF_SUBGRAPHS}.pt'
EXPERIMENT_RES_PATH = f'neighbor_deeper_num_{NUMBER_OF_SUBGRAPHS}_experiment_res.pk'
args = ArgsInit().args
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygNodePropPredDataset(name=args.dataset)
graph = dataset[0]
print(graph)
num_parts = NUMBER_OF_SUBGRAPHS
data_list = list(
NeighborSubgraphLoader(graph, num_parts=NUMBER_OF_SUBGRAPHS))
print(f'len of datalist: {len(data_list)}')
number_of_train = int(0.9 * num_parts)
train_data_list = data_list[0:number_of_train]
test_data_list = data_list[number_of_train:]
print(
f'Train test split successful, number of train: {len(train_data_list)} | number of test: {len(test_data_list)}'
)
args.in_channels = graph.x.size(-1)
args.num_tasks = dataset.num_classes
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
start_time = time.time()
epoch_loss_list = []
epoch_acc_list = []
highest_acc = 0
best_model_dict = None
for epoch in range(1, EPOCHS + 1):
epoch_loss, epoch_acc = train(data_list, model, optimizer, device)
epoch_loss_list.append(epoch_loss)
epoch_acc_list.append(epoch_acc)
print('Epoch {}, training loss {:.4f} | training acc {}'.format(
epoch, epoch_loss, epoch_acc))
test_acc = test(model, test_data_list)
if test_acc > highest_acc:
highest_acc = test_acc
best_model_dict = model.state_dict()
logging.info(
f'best test acc: {highest_acc} | saved to path {CKPT_PATH}')
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
experiment_result = {}
experiment_result['Total training time'] = total_time
experiment_result['Epoch loss list'] = epoch_loss_list
experiment_result['Epoch acc list'] = epoch_acc_list
experiment_result['Best test acc'] = highest_acc
torch.save(best_model_dict, CKPT_PATH)
with open(EXPERIMENT_RES_PATH, 'wb') as f:
pk.dump(experiment_result, f)
def main_get_mask(args, imp_num, rewind_weight_mask=None, rewind_predict_weight=None, resume_train_ckpt=None):
device = torch.device("cuda:" + str(args.device))
dataset = PygLinkPropPredDataset(name=args.dataset)
data = dataset[0]
# Data(edge_index=[2, 2358104], edge_weight=[2358104, 1], edge_year=[2358104, 1], x=[235868, 128])
split_edge = dataset.get_edge_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
edge_index = data.edge_index.to(device)
args.in_channels = data.x.size(-1)
args.num_tasks = 1
model = DeeperGCN(args).to(device)
pruning.add_mask(model, args)
predictor = LinkPredictor(args).to(device)
pruning.add_trainable_mask_noise(model, args, c=1e-4)
optimizer = torch.optim.Adam(list(model.parameters()) + list(predictor.parameters()), lr=args.lr)
results = {'epoch': 0 }
keys = ['highest_valid', 'final_train', 'final_test', 'highest_train']
hits = ['Hits@10', 'Hits@50', 'Hits@100']
for key in keys:
results[key] = {k: 0 for k in hits}
start_epoch = 1
if resume_train_ckpt:
start_epoch = resume_train_ckpt['epoch']
rewind_weight_mask = resume_train_ckpt['rewind_weight_mask']
ori_model_dict = model.state_dict()
over_lap = {k : v for k, v in resume_train_ckpt['model_state_dict'].items() if k in ori_model_dict.keys()}
ori_model_dict.update(over_lap)
model.load_state_dict(ori_model_dict)
print("Resume at IMP:[{}] epoch:[{}] len:[{}/{}]!".format(imp_num, resume_train_ckpt['epoch'], len(over_lap.keys()), len(ori_model_dict.keys())))
optimizer.load_state_dict(resume_train_ckpt['optimizer_state_dict'])
adj_spar, wei_spar = pruning.print_sparsity(model, args)
else:
rewind_weight_mask = copy.deepcopy(model.state_dict())
rewind_predict_weight = copy.deepcopy(predictor.state_dict())
for epoch in range(start_epoch, args.mask_epochs + 1):
t0 = time.time()
epoch_loss, prune_info_dict = train.train_mask(model, predictor,
x,
edge_index,
split_edge,
optimizer, args)
result = train.test(model, predictor,
x,
edge_index,
split_edge,
evaluator,
args.batch_size, args)
k = 'Hits@50'
train_result, valid_result, test_result = result[k]
if train_result > results['highest_train'][k]:
results['highest_train'][k] = train_result
if valid_result > results['highest_valid'][k]:
results['highest_valid'][k] = valid_result
results['final_train'][k] = train_result
results['final_test'][k] = test_result
results['epoch'] = epoch
pruning.save_all(model, predictor,
rewind_weight_mask,
optimizer,
imp_num,
epoch,
args.model_save_path,
'IMP{}_train_ckpt'.format(imp_num))
epoch_time = (time.time() - t0) / 60
print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'IMP:[{}] (GET Mask) Epoch:[{}/{}] LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}] | Adj[{:.3f}%] Wei[{:.3f}%] Time:[{:.2f}min]'
.format(imp_num, epoch, args.mask_epochs,
epoch_loss,
train_result * 100,
valid_result * 100,
test_result * 100,
results['final_test'][k] * 100,
results['epoch'],
prune_info_dict['adj_spar'],
prune_info_dict['wei_spar'],
epoch_time))
rewind_weight_mask, adj_spar, wei_spar = pruning.change(rewind_weight_mask, model, args)
print('-' * 100)
print("INFO : IMP:[{}] (GET MASK) Final Result Train:[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Adj:[{:.3f}%] Wei:[{:.3f}%]"
.format(imp_num, results['final_train'][k] * 100,
results['highest_valid'][k] * 100,
results['final_test'][k] * 100,
adj_spar,
wei_spar))
print('-' * 100)
return rewind_weight_mask, rewind_predict_weight
def main():
args = ArgsInit().save_exp()
if args.use_gpu:
device = torch.device("cuda:" +
str(args.device)) if torch.cuda.is_available(
) else torch.device("cpu")
else:
device = torch.device('cpu')
dataset = PygNodePropPredDataset(name=args.dataset)
data = dataset[0]
split_idx = dataset.get_idx_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
y_true = data.y.to(device)
train_idx = split_idx['train'].to(device)
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
if args.self_loop:
edge_index = add_self_loops(edge_index, num_nodes=data.num_nodes)[0]
sub_dir = 'SL_{}'.format(args.self_loop)
args.in_channels = data.x.size(-1)
args.num_tasks = dataset.num_classes
logging.info('%s' % args)
model = DeeperGCN(args).to(device)
logging.info(model)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
epoch_loss = train(model, x, edge_index, y_true, train_idx, optimizer)
logging.info('Epoch {}, training loss {:.4f}'.format(
epoch, epoch_loss))
model.print_params(epoch=epoch)
result = test(model, x, edge_index, y_true, split_idx, evaluator)
logging.info(result)
train_accuracy, valid_accuracy, test_accuracy = result
if train_accuracy > results['highest_train']:
results['highest_train'] = train_accuracy
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
sub_dir,
name_post='valid_best')
logging.info("%s" % results)
end_time = time.time()
total_time = end_time - start_time
logging.info('Total time: {}'.format(
time.strftime('%H:%M:%S', time.gmtime(total_time))))
# We will learn to deform the source mesh by offsetting its vertices
# The shape of the deform parameters is equal to the total number of vertices in src_mesh#
#deform_verts = torch.full(src_mesh.verts_packed().shape, 0.0, device=device, requires_grad=True)
#model = GCN()
print(device)
#TODO: more features
gcn1 = DeeperGCN(3, 128, 3, 12)
gcn1.to(device)
gcn2 = DeeperGCN(128 + 3, 128, 3, 12)
gcn2.to(device)
gcn3 = DeeperGCN(128 + 3, 128, 3, 12)
gcn3.to(device)
# The optimizer
optimizer = torch.optim.Adam(list(gcn1.parameters()) +
list(gcn2.parameters()) +
list(gcn3.parameters()), lr=0.0001)
#optimizer = torch.optim.SGD([deform_verts], lr=1.0, momentum=0.9)
# Number of optimization steps
Niter = 5001
# Weight for the chamfer loss
w_chamfer = 1.0
# Weight for mesh edge loss
w_edge = 0.1#1.0
# Weight for mesh normal consistency
w_normal = 1.6e-4#0.01
# Weight for mesh laplacian smoothing
w_laplacian = 0.3#0.1
def main(args):
device = torch.device("cuda:" + str(args.device))
dataset = OGBNDataset(dataset_name=args.dataset)
nf_path = dataset.extract_node_features(args.aggr)
args.num_tasks = dataset.num_tasks
args.nf_path = nf_path
evaluator = Evaluator(args.dataset)
criterion = torch.nn.BCEWithLogitsLoss()
valid_data_list = []
for i in range(args.num_evals):
parts = dataset.random_partition_graph(
dataset.total_no_of_nodes,
cluster_number=args.valid_cluster_number)
valid_data = dataset.generate_sub_graphs(
parts, cluster_number=args.valid_cluster_number)
valid_data_list.append(valid_data)
model = DeeperGCN(args).to(device)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0,
'epoch': 0
}
start_time = time.time()
for epoch in range(1, args.epochs + 1):
# do random partition every epoch
train_parts = dataset.random_partition_graph(
dataset.total_no_of_nodes, cluster_number=args.cluster_number)
data = dataset.generate_sub_graphs(train_parts,
cluster_number=args.cluster_number)
epoch_loss = train.train_baseline(data, dataset, model, optimizer,
criterion, device, args)
result = train.multi_evaluate(valid_data_list, dataset, model,
evaluator, device)
train_result = result['train']['rocauc']
valid_result = result['valid']['rocauc']
test_result = result['test']['rocauc']
if valid_result > results['highest_valid']:
results['highest_valid'] = valid_result
results['final_train'] = train_result
results['final_test'] = test_result
results['epoch'] = epoch
save_dir = save_ckpt(model,
optimizer,
round(epoch_loss, 4),
epoch,
args.model_save_path,
name_post='valid_best')
print(
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'(Baseline) Epoch:[{}/{}] LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}]'
.format(epoch, args.epochs, epoch_loss, train_result *
100, valid_result * 100, test_result *
100, results['final_test'] * 100, results['epoch']))
end_time = time.time()
total_time = end_time - start_time
# logging.info('Total time: {}'.format(time.strftime('%H:%M:%S', time.gmtime(total_time))))
print("=" * 120)
print(
"syd final: BASELINE, Train:[{:.2f}] Best Val:[{:.2f}] at epoch:[{}] | Final Test Acc:[{:.2f}] Total time:[{}]"
.format(results['final_train'] * 100, results['highest_valid'] * 100,
results['epoch'], results['final_test'] * 100,
time.strftime('%H:%M:%S', time.gmtime(total_time))))
print("=" * 120)
def main_fixed_mask(args, imp_num, resume_train_ckpt=None):
device = torch.device("cuda:" + str(args.device))
dataset = PygLinkPropPredDataset(name=args.dataset)
data = dataset[0]
# Data(edge_index=[2, 2358104], edge_weight=[2358104, 1], edge_year=[2358104, 1], x=[235868, 128])
split_edge = dataset.get_edge_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
edge_index = data.edge_index.to(device)
args.in_channels = data.x.size(-1)
args.num_tasks = 1
model = DeeperGCN(args).to(device)
pruning.add_mask(model, args)
predictor = LinkPredictor(args).to(device)
rewind_weight_mask, adj_spar, wei_spar = pruning.resume_change(resume_train_ckpt, model, args)
model.load_state_dict(rewind_weight_mask)
predictor.load_state_dict(resume_train_ckpt['predictor_state_dict'])
# model.load_state_dict(rewind_weight_mask)
# predictor.load_state_dict(rewind_predict_weight)
adj_spar, wei_spar = pruning.print_sparsity(model, args)
for name, param in model.named_parameters():
if 'mask' in name:
param.requires_grad = False
optimizer = torch.optim.Adam(list(model.parameters()) + list(predictor.parameters()), lr=args.lr)
#results = {}
results = {'epoch': 0 }
keys = ['highest_valid', 'final_train', 'final_test', 'highest_train']
hits = ['Hits@10', 'Hits@50', 'Hits@100']
for key in keys:
results[key] = {k: 0 for k in hits}
results['adj_spar'] = adj_spar
results['wei_spar'] = wei_spar
start_epoch = 1
for epoch in range(start_epoch, args.fix_epochs + 1):
t0 = time.time()
epoch_loss = train.train_fixed(model, predictor, x, edge_index, split_edge, optimizer, args.batch_size, args)
result = train.test(model, predictor, x, edge_index, split_edge, evaluator, args.batch_size, args)
# return a tuple
k = 'Hits@50'
train_result, valid_result, test_result = result[k]
if train_result > results['highest_train'][k]:
results['highest_train'][k] = train_result
if valid_result > results['highest_valid'][k]:
results['highest_valid'][k] = valid_result
results['final_train'][k] = train_result
results['final_test'][k] = test_result
results['epoch'] = epoch
pruning.save_all(model, predictor,
rewind_weight_mask,
optimizer,
imp_num,
epoch,
args.model_save_path,
'IMP{}_fixed_ckpt'.format(imp_num))
epoch_time = (time.time() - t0) / 60
print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'IMP:[{}] (FIX Mask) Epoch:[{}/{}] LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}] Time:[{:.2f}min]'
.format(imp_num, epoch, args.fix_epochs, epoch_loss, train_result * 100,
valid_result * 100,
test_result * 100,
results['final_test'][k] * 100,
results['epoch'],
epoch_time))
print("=" * 120)
print("syd final: IMP:[{}], Train:[{:.2f}] Best Val:[{:.2f}] at epoch:[{}] | Final Test Acc:[{:.2f}] Adj:[{:.2f}%] Wei:[{:.2f}%]"
.format(imp_num, results['final_train'][k] * 100,
results['highest_valid'][k] * 100,
results['epoch'],
results['final_test'][k] * 100,
results['adj_spar'],
results['wei_spar']))
print("=" * 120)
def main_get_mask(args,
imp_num,
rewind_weight_mask=None,
resume_train_ckpt=None):
device = torch.device("cuda:" + str(args.device))
dataset = PygNodePropPredDataset(name=args.dataset)
data = dataset[0]
split_idx = dataset.get_idx_split()
evaluator = Evaluator(args.dataset)
x = data.x.to(device)
y_true = data.y.to(device)
train_idx = split_idx['train'].to(device)
edge_index = data.edge_index.to(device)
edge_index = to_undirected(edge_index, data.num_nodes)
if args.self_loop:
edge_index = add_self_loops(edge_index, num_nodes=data.num_nodes)[0]
args.in_channels = data.x.size(-1)
args.num_tasks = dataset.num_classes
print("-" * 120)
model = DeeperGCN(args).to(device)
pruning.add_mask(model, args.num_layers)
if rewind_weight_mask:
model.load_state_dict(rewind_weight_mask)
adj_spar, wei_spar = pruning.print_sparsity(model, args)
pruning.add_trainable_mask_noise(model, args, c=1e-5)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
results = {
'highest_valid': 0,
'final_train': 0,
'final_test': 0,
'highest_train': 0,
'epoch': 0
}
start_epoch = 1
if resume_train_ckpt:
start_epoch = resume_train_ckpt['epoch']
rewind_weight_mask = resume_train_ckpt['rewind_weight_mask']
ori_model_dict = model.state_dict()
over_lap = {
k: v
for k, v in resume_train_ckpt['model_state_dict'].items()
if k in ori_model_dict.keys()
}
ori_model_dict.update(over_lap)
model.load_state_dict(ori_model_dict)
print("Resume at IMP:[{}] epoch:[{}] len:[{}/{}]!".format(
imp_num, resume_train_ckpt['epoch'], len(over_lap.keys()),
len(ori_model_dict.keys())))
optimizer.load_state_dict(resume_train_ckpt['optimizer_state_dict'])
adj_spar, wei_spar = pruning.print_sparsity(model, args)
else:
rewind_weight_mask = copy.deepcopy(model.state_dict())
for epoch in range(start_epoch, args.mask_epochs + 1):
epoch_loss = train(model, x, edge_index, y_true, train_idx, optimizer,
args)
result = test(model, x, edge_index, y_true, split_idx, evaluator)
train_accuracy, valid_accuracy, test_accuracy = result
if valid_accuracy > results['highest_valid']:
results['highest_valid'] = valid_accuracy
results['final_train'] = train_accuracy
results['final_test'] = test_accuracy
results['epoch'] = epoch
rewind_weight_mask, adj_spar, wei_spar = pruning.get_final_mask_epoch(
model, rewind_weight_mask, args)
#pruning.save_all(model, rewind_weight_mask, optimizer, imp_num, epoch, args.model_save_path, 'IMP{}_train_ckpt'.format(imp_num))
print(
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + ' | ' +
'IMP:[{}] (GET Mask) Epoch:[{}/{}]\t LOSS:[{:.4f}] Train :[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Update Test:[{:.2f}] at epoch:[{}] | Adj:[{:.2f}%] Wei:[{:.2f}%]'
.format(imp_num, epoch, args.mask_epochs, epoch_loss,
train_accuracy * 100, valid_accuracy * 100, test_accuracy *
100, results['final_test'] *
100, results['epoch'], adj_spar, wei_spar))
print('-' * 100)
print(
"INFO : IMP:[{}] (GET MASK) Final Result Train:[{:.2f}] Valid:[{:.2f}] Test:[{:.2f}] | Adj:[{:.2f}%] Wei:[{:.2f}%] "
.format(imp_num, results['final_train'] * 100,
results['highest_valid'] * 100, results['final_test'] * 100,
adj_spar, wei_spar))
print('-' * 100)
return rewind_weight_mask