GNN을 활용한 노드 분류

노드 분류 과제

Cora 데이터셋

# Cora:
#   nodes  = 2708 papers
#   edges  = citations
#   features per node = 1433
#   classes = 7

Planetoid로 불러오기

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root="data/Cora", name="Cora")
data = dataset[0]
print(data)  # x, edge_index, y, train_mask, test_mask

데이터 객체

7개 클래스를 위한 GCN

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(1433, 64)
        self.conv2 = GCNConv(64, 7)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, training=self.training)
        return self.conv2(x, edge_index)

전이적 학습

훈련 마스크

model = GCN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

마스크된 손실의 직관

평가 마스크로 평가하기

@torch.no_grad()
def test():
    model.eval()
    pred = model(data.x, data.edge_index).argmax(dim=1)
    correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
    return int(correct) / int(data.test_mask.sum())

t-SNE로 임베딩 시각화하기

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

emb = model.conv1(data.x, data.edge_index).detach().numpy()
z = TSNE(n_components=2).fit_transform(emb)
plt.scatter(z[:, 0], z[:, 1], c=data.y, cmap="tab10", s=8)

t-SNE 그래프 해석하기

빠른 확인

복습