def main():
saved_graph = os.path.join('assets', 'saved_graphs', 'best_dgi.pickle')
saved_logreg = os.path.join('assets', 'saved_graphs', 'best_logreg.pickle')
dataset = 'cora'
# training params
batch_size = 1
nb_epochs = 10000
patience = 25
lr = 0.001
l2_coef = 0.0
drop_prob = 0.0
hid_units = 512
sparse = True
nonlinearity = 'prelu' # special name to separate parameters
adj, features, labels, idx_train, idx_test, idx_val = process.load_data(dataset)
features, _ = process.preprocess_features(features)
nb_nodes = features.shape[0]
ft_size = features.shape[1]
nb_classes = labels.shape[1]
adj = process.normalize_adj(adj + sp.eye(adj.shape[0]))
if sparse:
adj = process.sparse_mx_to_torch_sparse_tensor(adj)
else:
adj = (adj + sp.eye(adj.shape[0])).todense()
features = torch.FloatTensor(features[np.newaxis])
if not sparse:
adj = torch.FloatTensor(adj[np.newaxis])
labels = torch.FloatTensor(labels[np.newaxis])
idx_train = torch.LongTensor(idx_train)
idx_val = torch.LongTensor(idx_val)
idx_test = torch.LongTensor(idx_test)
print("Training Nodes: {}, Testing Nodes: {}, Validation Nodes: {}".format(len(idx_train), len(idx_test), len(idx_val)))
model = DGI(ft_size, hid_units, nonlinearity)
optimiser = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2_coef)
if torch.cuda.is_available():
print('Using CUDA')
model.cuda()
features = features.cuda()
if sparse:
sp_adj = sp_adj.cuda()
else:
adj = adj.cuda()
labels = labels.cuda()
idx_train = idx_train.cuda()
idx_val = idx_val.cuda()
idx_test = idx_test.cuda()
b_xent = nn.BCEWithLogitsLoss()
xent = nn.CrossEntropyLoss()
cant_wait = 0
best = 1e9
best_t = 0
if not os.path.exists(saved_graph):
pbar = trange(nb_epochs)
for epoch in pbar:
model.train()
optimiser.zero_grad()
idx = np.random.permutation(nb_nodes)
shuf_fts = features[:, idx, :]
lbl_1 = torch.ones(batch_size, nb_nodes)
lbl_2 = torch.zeros(batch_size, nb_nodes)
lbl = torch.cat((lbl_1, lbl_2), 1)
if torch.cuda.is_available():
shuf_fts = shuf_fts.cuda()
lbl = lbl.cuda()
logits = model(features, shuf_fts, adj, sparse, None, None, None)
loss = b_xent(logits, lbl)
pbar.desc = 'Loss: {:.4f}'.format(loss)
if loss < best:
best = loss
best_t = epoch
cnt_wait = 0
torch.save(model.state_dict(), saved_graph)
else:
cant_wait += 1
if cant_wait == patience:
tqdm.write('Early stopping!')
break
loss.backward()
optimiser.step()
print('Loading {}th Epoch'.format(best_t) if best_t else 'Loading Existing Graph')
model.load_state_dict(torch.load(saved_graph))
embeds, _ = model.embed(features, adj, sparse, None)
train_embs = embeds[0, idx_train]
val_embs = embeds[0, idx_val]
test_embs = embeds[0, idx_test]
train_lbls = torch.argmax(labels[0, idx_train], dim=1)
val_lbls = torch.argmax(labels[0, idx_val], dim=1)
test_lbls = torch.argmax(labels[0, idx_test], dim=1)
tot = torch.zeros(1)
if torch.cuda.is_available():
tot = tot.cuda()
accs = []
print("\nValidation:")
pbar = trange(50)
for _ in pbar:
log = LogReg(hid_units, nb_classes)
opt = torch.optim.Adam(log.parameters(), lr=0.01, weight_decay=0.0)
pat_steps = 0
best_acc = torch.zeros(1)
if torch.cuda.is_available():
log.cuda()
best_acc = best_acc.cuda()
for _ in range(100):
log.train()
opt.zero_grad()
logits = log(train_embs)
loss = xent(logits, train_lbls)
loss.backward()
opt.step()
logits = log(test_embs)
preds = torch.argmax(logits, dim=1)
acc = torch.sum(preds == test_lbls).float() / test_lbls.shape[0]
accs.append(acc * 100)
pbar.desc = "Accuracy: {:.2f}%".format(100 * acc)
tot += acc
torch.save(log.state_dict(), saved_logreg)
accs = torch.stack(accs)
print('Average Accuracy: {:.2f}%'.format(accs.mean()))
print('Standard Deviation: {:.3f}'.format(accs.std()))
print("\nTesting")
logits = log(val_embs)
preds = torch.argmax(logits, dim=1)
acc = torch.sum(preds == val_lbls).float() / val_lbls.shape[0]
print("Accuracy: {:.2f}%".format(100 * acc))
opt = torch.optim.Adam(log.parameters(), lr=0.01, weight_decay=0.0)
if args.cuda:
log.cuda()
for epoch in range(100000):
log.train()
opt.zero_grad()
logits = log(train_embs)
loss = xent(logits, train_lbls)
logits_val = log(val_embs)
loss_val = xent(logits_val, val_lbls)
loss_values.append(loss_val)
# print("train_loss: "+ str(loss) +" "+"val_loss: "+ str(loss_val) )
loss.backward()
opt.step()
torch.save(log.state_dict(), '{}.mlp.pkl'.format(epoch))
if loss_values[-1] < best:
best = loss_values[-1]
best_epoch = epoch
bad_counter = 0
else:
bad_counter += 1
if bad_counter == train_patience:
break
files = glob.glob('*.mlp.pkl')
for file in files:
epoch_nb = int(file.split('.')[0])
if epoch_nb < best_epoch:
os.remove(file)
