torch_geometric_gcn_node_classification_cora_dataset.py

1import torch
2import torch.nn.functional as F
3from torch_geometric.nn import GCNConv
4from torch_geometric.datasets import Planetoid
5
6# 1. Load the Cora dataset
7dataset = Planetoid(root='/tmp/Cora', name='Cora')
8data = dataset[0]
9
10# 2. Define the GCN architecture
11class GCN(torch.nn.Module):
12    def __init__(self):
13        super().__init__()
14        self.conv1 = GCNConv(dataset.num_node_features, 16)
15        self.conv2 = GCNConv(16, dataset.num_classes)
16
17    def forward(self, data):
18        x, edge_index = data.x, data.edge_index
19
20        x = self.conv1(x, edge_index)
21        x = F.relu(x)
22        x = F.dropout(x, training=self.training)
23        x = self.conv2(x, edge_index)
24
25        return F.log_softmax(x, dim=1)
26
27# 3. Initialize model, optimizer, and training loop
28device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
29model = GCN().to(device)
30data = data.to(device)
31optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
32
33model.train()
34for epoch in range(200):
35    optimizer.zero_grad()
36    out = model(data)
37    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
38    loss.backward()
39    optimizer.step()
40
41# 4. Evaluate the model
42model.eval()
43pred = model(data).argmax(dim=1)
44correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
45acc = int(correct) / int(data.test_mask.sum())
46print(f'Accuracy: {acc:.4f}')