pytorch_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 dataset
7dataset = Planetoid(root='/tmp/Cora', name='Cora')
8
9# 2. Define the GCN architecture
10class GCN(torch.nn.Module):
11    def __init__(self):
12        super().__init__()
13        self.conv1 = GCNConv(dataset.num_node_features, 16)
14        self.conv2 = GCNConv(16, dataset.num_classes)
15
16    def forward(self, data):
17        x, edge_index = data.x, data.edge_index
18
19        x = self.conv1(x, edge_index)
20        x = F.relu(x)
21        x = F.dropout(x, training=self.training)
22        x = self.conv2(x, edge_index)
23
24        return F.log_softmax(x, dim=1)
25
26# 3. Initialize model, optimizer, and data
27device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
28model = GCN().to(device)
29data = dataset[0].to(device)
30optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
31
32# 4. Training loop
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# 5. Evaluation
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}')