adj_feat_list = load_list(args.dataset, args.K)
partition = {"train": adj_feat_list}
# Parameters
train_params = {'dim': dim,
'adj_dim': (dim, dim),
'feature_dim': (dim, 1),
'batch_size': batch_size,
'shuffle': False,
'path': path}
training_generator = DataGenerator(partition['train'], **train_params)
print("the training data is ready")
model = GCNModelVAE(batch_size, args.hidden1, args.hidden2, args.dropout)
model.cuda() ###################################################################
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# scheduler = ReduceLROnPlateau(optimizer, 'min')
scheduler = MultiStepLR(optimizer, milestones=[50, 80], gamma=0.1)
loss_record = []
start_time = time.time()
for epoch in range(epochs):
scheduler.step()
# Training
# numm=0
cur_loss_list = []
for adj_bs, features_bs in training_generator:
# Transfer to GPU
# print(int(adj_bs.shape[1]/dim))
adj_label = adj_bs + sp.eye(adj_bs.shape[0])
# adj_label = sparse_to_tuple(adj_label)
adj_label = torch.FloatTensor(adj_label.toarray())
# adj_label = torch.LongTensor(adj_label.toarray())
pos_weight = float(adj_bs.shape[0] * adj_bs.shape[0] -
adj_bs.sum()) / adj_bs.sum()
norm = adj_bs.shape[0] * adj_bs.shape[0] / float(
(adj_bs.shape[0] * adj_bs.shape[0] - adj_bs.sum()) * 2)
model = GCNModelVAE(feat_dim, args.hidden1, args.hidden2,
args.dropout)
model = torch.nn.DataParallel(
model, device_ids=[0]
) ###################################################################
model.cuda()
# parallel_model = DataParallelModel(model, device_ids=[0])
# parallel_loss = DataParallelCriterion(loss_function, device_ids=[0])
# optimizer = optim.Adam(model.parameters(), lr=args.lr)
# parallel_model.train()
# optimizer.zero_grad()
# recovered, mu, logvar = parallel_model(features_bs, adj_norm)
# loss = loss_function(preds=recovered,
# labels=adj_label,
# mu=mu,
# logvar=logvar,
# n_nodes=n_nodes,
# norm=norm,
# pos_weight=pos_weight)
#
# cpu_loss = loss.cpu()
def gae_for(args):
print("Using {} dataset".format(args.dataset_str))
#得到原始邻接矩阵和特征
adj, features = load_data(args.dataset_str)
n_nodes, feat_dim = features.shape
# Store original adjacency matrix (without diagonal entries) for later
adj_orig = adj
adj_orig = adj_orig - sp.dia_matrix(
(adj_orig.diagonal()[np.newaxis, :], [0]), shape=adj_orig.shape)
adj_orig.eliminate_zeros()
print(adj.shape)
adj_train, train_edges, val_edges, val_edges_false, test_edges, test_edges_false = mask_test_edges(
adj)
adj = adj_train
print('adj.shape:{}'.format(adj.shape))
print("np.sum(adj):{}".format(np.sum(adj)))
print("len(train_edges):{}".format(len(train_edges)))
print("len(test_edges):{}".format(len(test_edges)))
print("len(val_edges):{}".format(len(val_edges)))
# Some preprocessing
adj_norm = preprocess_graph(adj)
adj_label = adj_train + sp.eye(adj_train.shape[0])
# adj_label = sparse_to_tuple(adj_label)
adj_label = torch.FloatTensor(adj_label.toarray())
pos_weight = float(adj.shape[0] * adj.shape[0] - adj.sum()) / adj.sum()
# pos_weight=torch.from_numpy(pos_weight)
norm = adj.shape[0] * adj.shape[0] / float(
(adj.shape[0] * adj.shape[0] - adj.sum()) * 2)
model = GCNModelVAE(feat_dim, args.hidden1, args.hidden2, args.dropout)
if cuda:
model = model.cuda()
features = features.cuda()
adj_norm = adj_norm.cuda()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
hidden_emb = None
for epoch in range(args.epochs):
t = time.time()
model.train()
optimizer.zero_grad()
#每次将整个图放入到模型中训练
recovered, mu, logvar = model(features, adj_norm)
loss = loss_function(preds=recovered,
labels=adj_label,
mu=mu,
logvar=logvar,
n_nodes=n_nodes,
norm=norm)
loss.backward()
cur_loss = loss.item()
optimizer.step()
hidden_emb = mu.cpu().data.numpy()
roc_curr, ap_curr = get_roc_score(hidden_emb, adj_orig, val_edges,
val_edges_false)
print("Epoch:", '%04d' % (epoch + 1), "train_loss=",
"{:.5f}".format(cur_loss), "val_ap=", "{:.5f}".format(ap_curr),
"time=", "{:.5f}".format(time.time() - t))
print("Optimization Finished!")
roc_score, ap_score = get_roc_score(hidden_emb, adj_orig, test_edges,
test_edges_false)
print('Test ROC score: ' + str(roc_score))
print('Test AP score: ' + str(ap_score))
model = model.cpu()
