#Load data
adj, features, labels, idx_train, idx_val, idx_test = load_data()
model = GAT(nfeat=features.shape[1],
nhid=args.hidden,
nclass=int(labels.max())+1,
dropout=args.dropout,
nheads=args.nb_heads,
alpha=args.alpha)
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
if args.cuda:
model.cuda()
features = features.cuda()
adj = adj.cuda()
labels = labels.cuda()
idx_train = idx_train.cuda()
idx_val = idx_val.cuda()
idx_test = idx_test.cuda()
features, adj, labels = Variable(features), Variable(adj), Variable(labels)
train_losses, train_accs, val_losses, val_accs = [], [], [], []
def train(epoch):
t = time.time()
model.train()
optimizer.zero_grad()
print("Let's use multi-GPUs!")
model = nn.DataParallel(model)
else:
print("Only use one GPU")
labels = adj.view(node_num * node_num).type_as(idx_train)
labels = labels.type(torch.FloatTensor)
ori_labels = ori_adj.view(node_num * node_num).type_as(idx_train)
ori_labels = ori_labels.type(torch.FloatTensor)
criterion = nn.BCELoss()
adj, ori_labels = Variable(adj), Variable(ori_labels)
features = Variable(features)
labels = Variable(labels)
if args.cuda:
model = model.cuda()
features = features.cuda()
adj = adj.cuda()
labels = labels.cuda()
idx_train = idx_train.cuda()
idx_val = idx_val.cuda()
idx_test = idx_test.cuda()
criterion = criterion.cuda()
ori_labels = ori_labels.cuda()
writer = SummaryWriter(
comment='ori_0.1_hidden_{}_lr_{}_#batch_{}_SampleSize_{}_testBatchSize_{}'.
format(args.hidden, args.lr, args.batchSize, args.sampleSize,
args.testBatchSize))
train_batches = iterate_return(idx_train, args.sampleSize, labels, adj,
def main(opt):
data = get_dataset(opt)
g = data[0]
if opt['gpu'] < 0:
cuda = False
else:
cuda = True
g = g.int().to(opt['gpu'])
features = g.ndata['feat']
labels = g.ndata['label']
train_mask = g.ndata['train_mask']
val_mask = g.ndata['val_mask']
test_mask = g.ndata['test_mask']
num_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes, train_mask.int().sum().item(),
val_mask.int().sum().item(), test_mask.int().sum().item()))
# add self loop
g = dgl.remove_self_loop(g)
g = dgl.add_self_loop(g)
n_edges = g.number_of_edges()
# create model
heads = ([opt['num_heads']] * opt['num_layers']) + [opt['num_out_heads']]
if opt['model'] == 'GAT':
model = GAT(g, opt['num_layers'], num_feats, opt['num_hidden'],
n_classes, heads, F.elu, opt['in_drop'], opt['attn_drop'],
opt['negative_slope'], opt['residual'], opt)
elif opt['model'] == 'AGNN':
model = AGNN(g, opt['num_layers'], num_feats, opt['num_hidden'],
n_classes, opt['in_drop'], opt)
print(model)
if opt['early_stop']:
stopper = EarlyStopping(patience=100)
if cuda:
model.cuda()
# use optimizer
optimizer = get_optimizer(opt['optimizer'],
parameters=model.parameters(),
lr=opt['lr'],
weight_decay=opt['weight_decay'])
# initialize graph
dur = []
for epoch in range(opt['epochs']):
# model.train()
if epoch >= 3:
t0 = time.time()
# forward
# logits = model(features)
# loss = loss_fcn(logits[train_mask], labels[train_mask])
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
loss, logits = train(model, optimizer, features, train_mask, labels)
if epoch >= 3:
dur.append(time.time() - t0)
train_acc = accuracy(logits[train_mask], labels[train_mask])
if opt['fastmode']:
val_acc = accuracy(logits[val_mask], labels[val_mask])
else:
val_acc = evaluate(model, features, labels, val_mask)
if opt['early_stop']:
if stopper.step(val_acc, model):
break
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | TrainAcc {:.4f} |"
" ValAcc {:.4f} | ETputs(KTEPS) {:.2f}".format(
epoch, np.mean(dur), loss.item(), train_acc, val_acc,
n_edges / np.mean(dur) / 1000))
print()
if opt['early_stop']:
model.load_state_dict(torch.load('es_checkpoint.pt'))
acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# Load data
adj, features, labels, idx_train, idx_val, idx_test = load_data()
# Model and optimizer
model = GAT(nfeat=features.shape[1], nhid=args.hidden, nclass=int(labels.max()) + 1, dropout=args.dropout, nheads=args.nb_heads, alpha=args.alpha)
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
if args.cuda:
model.cuda()
features = features.cuda()
adj = adj.cuda()
labels = labels.cuda()
idx_train = idx_train.cuda()
idx_val = idx_val.cuda()
idx_test = idx_test.cuda()
features, adj, labels = Variable(features), Variable(adj), Variable(labels)
def train(epoch):
t = time.time()
model.train()
optimizer.zero_grad()
output = model(features, adj)
