Train DGL GCN Model

dgl-train-karate.py

optimizer_name = "Adam"
lr = 1e-1
optimizer = getattr(torch.optim, optimizer_name)(model.parameters(), lr=lr)
epochs = 100

def train():
  model.train()
  optimizer.zero_grad()
  F.nll_loss(model(data, node_features)[train_mask],
             node_labels[train_mask]).backward()
  optimizer.step()

@torch.no_grad()
def test():
  model.eval()
  logits = model(data, node_features)
  mask1 = train_mask
  pred1 = logits[mask1].max(1)[1]
  acc1 = pred1.eq(node_labels[mask1]).sum().item() / mask1.sum().item()
  mask = test_mask
  pred = logits[mask].max(1)[1]
  acc = pred.eq(node_labels[mask]).sum().item() / mask.sum().item()
  return acc1,acc

for epoch in range(1, epochs):
  train()
  train_acc,test_acc = test()

print('#' * 50)
print('Train Accuracy: %s' %train_acc )
print('Test Accuracy: %s' % test_acc)
print('#' * 50)