Python) Cora dataset을 활용하여 Link Prediction

Objective

GNN을 사용하여 어떻게 Link Prediction을 할 수 있는지에 대해서 알아보고자 함.

 

 

Introduction

많은 예제들이 보통 node 분류를 하는 데, GNN 많은 예제들이 있고, Link Prediction을 다루는 예제는 많이 없다.

개인적으로 Node 분류도 관심이 있지만, Edge를 어떻게 예측하는지 관심이 있어 시작하게 되었고 여기서는 어떻게 예측을 하는지 알아보고자 한다. 

 

DataSet

여기서는 자주 사용하는 Cora 데이터 세트로 진행하려고 한다. 

이 데이터의 Node와 Edge의 의미는 다음과 같다.

Node	머신러닝 논문
Edge	논문 쌍 간의 인용

즉 이 논문은 인용 그래프를 나타낸다고 할 수 있다. 

Node (Node Classification)	논문 분류 (7개의 클래스)
Edge (Link Prediction)	인용 여부

 

Implementation

기존 코드를 활용하고 본 결과를 말씀드리자면, 직접적으로 예측한다기보다는, 임베딩 값을 계산하고 그 임베딩 된 값(노드)들끼리의 곱해서 합을 낸 결과에 sigmoid를 통해서 확률 값을 계산하는 것으로 확인했다. 

 

Load Library 

import os.path as osp

import torch
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np
import random

from torch_geometric.utils import negative_sampling
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
from torch_geometric.nn import GCNConv
from torch_geometric.utils import train_test_split_edges

device = "cpu"

Load Data 

여기서는 기존 데이터를 train_test_split_edges를 통해서 edge에 따라서 데이터를 분리해주는 작업을 해야 한다.

속성	설명
x	node features (num nodes , num node features)
edge_index	edge feature (2 , num edges)
pos_edge_index	연결되어 있는 edge 들을 의미하는 것으로 모인다.(추정)

# load the Cora dataset
dataset = 'Cora'
path = osp.join('../', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]
print(dataset.data)

# Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])

# use train_test_split_edges to create neg and positive edges
data.train_mask = data.val_mask = data.test_mask = data.y = None
data = train_test_split_edges(data)
print(data)

# Data(x=[2708, 1433], val_pos_edge_index=[2, 263], test_pos_edge_index=[2, 527], train_pos_edge_index=[2, 8976], train_neg_adj_mask=[2708, 2708], val_neg_edge_index=[2, 263], test_neg_edge_index=[2, 527])

Define Model

link prediction을 할 때 여기서는 encode와 decode로 나눠서 진행한다.

encode

다른 것들과 동일하게 기존 데이터의 node와 edge를 사용한다.

결과값으로는 해당 데이터에 뉴럴 네트워크를 거치고 나서 임베딩 된 값을 내본다.

마치 auto encoder에서 encoder와 같은 부분이다.

decode

학습을 하기 위해서는 연결되어 있는 edge와 연결되어 있지 않는 edge가 필요하고, 

여기서 연결되어 있지 않는 edge를 계산하기 위해서 negative_sampling 방법을 이용해서 추출한다.

그래서 decode에서는 negative sampling을 통해서 연결되지 않은 edge(neg_edge_index)와 연결되어 있는 pos_edge_index를 가지고 와서 각 edge 별로의 해당하는 node들을 이용해서 dot product 하여 logits 값을 계산한다.

 

decode_all

여기서는 특정 edge 말고 전체 node의 edge 연결성을 확인하는 코드이다. 

 

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, 128)
        self.conv2 = GCNConv(128, 64)

    def encode(self, x, edge_index):
        x = self.conv1(x, edge_index) # convolution 1
        x = x.relu()
        return self.conv2(x, edge_index) # convolution 2

    def decode(self, z, pos_edge_index, neg_edge_index): # only pos and neg edges
        edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1) # concatenate pos and neg edges
        logits = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)  # dot product 
        return logits

    def decode_all(self, z): 
        prob_adj = z @ z.t() # get adj NxN
        return (prob_adj > 0).nonzero(as_tuple=False).t() # get predicted edge_list

 

Function Logit, Link

def get_link_logits(model, x , edge_index , neg_edge_index) :
    z = model.encode(x , edge_index) #encode
    link_logits = model.decode(z, edge_index, neg_edge_index) # decode
    return link_logits
    
def get_link_labels(pos_edge_index, neg_edge_index):
    # returns a tensor:
    # [1,1,1,1,...,0,0,0,0,0,..] with the number of ones is equel to the lenght of pos_edge_index
    # and the number of zeros is equal to the length of neg_edge_index
    E = pos_edge_index.size(1) + neg_edge_index.size(1)
    link_labels = torch.zeros(E, dtype=torch.float, device=device)
    link_labels[:pos_edge_index.size(1)] = 1.
    return link_labels

get_link_logits는 위에서 말한 encode와 decode를 사용하여 logit을 출력해주는 함수이다.

get_link_labels는 edge index를 사용하여, link label을 만드는 함수이다.

 

Logit

neg_edge_index = negative_sampling(
        edge_index=data.train_pos_edge_index, #positive edges
        num_nodes=data.num_nodes, # number of nodes
        num_neg_samples=data.train_pos_edge_index.size(1)) # number of neg_sample equal to number of pos_edges

print(data.train_pos_edge_index.shape) # torch.Size([2, 8976])
print(neg_edge_index.shape) # torch.Size([2, 8976])
link_logits = get_link_logits(model , data.x , data.train_pos_edge_index,neg_edge_index)
print(link_logits.size()) # torch.Size([17952])

link_logits.sigmoid()
# tensor([0.5464, 0.5480, 0.5567,  ..., 0.5421, 0.5441, 0.5430], grad_fn=<SigmoidBackward0>)

위와 같은 식으로 logit값을 구할 수 있다. 

그리고 logit 값에 sigmoid를 취해서 확률값으로 바꿔준다.

 

Train

def train():
    model.train()
    neg_edge_index = negative_sampling(
        edge_index=data.train_pos_edge_index, #positive edges
        num_nodes=data.num_nodes, # number of nodes
        num_neg_samples=data.train_pos_edge_index.size(1)) # number of neg_sample equal to number of pos_edges

    link_logits = get_link_logits(model , data.x , data.train_pos_edge_index,neg_edge_index)
    optimizer.zero_grad()
    link_labels = get_link_labels(data.train_pos_edge_index, neg_edge_index)
    loss = F.binary_cross_entropy_with_logits(link_logits, link_labels)
    loss.backward()
    optimizer.step()

    return loss
train()

이제 학습은 다음과 같이 이뤄진다. 

실제로 이 코드는 우리가 기존에 알던 코드와 매우 유사하다.

 

 

Test

@torch.no_grad()
def test():
    model.eval()
    perfs = []
    for prefix in ["val", "test"]:
        pos_edge_index = data[f'{prefix}_pos_edge_index']
        neg_edge_index = data[f'{prefix}_neg_edge_index']
        
        link_logits = get_link_logits(model , data.x , pos_edge_index,neg_edge_index)
        
        link_probs = link_logits.sigmoid() # apply sigmoid
        
        link_labels = get_link_labels(pos_edge_index, neg_edge_index) # get link
        
        perfs.append(roc_auc_score(link_labels.cpu(), link_probs.cpu())) #compute roc_auc score
    return perfs

이런 식으로 연결 여부를 가지고 검증할 수 있다.

 

Node Embedding

dataset = Planetoid(path, "Cora", transform=T.NormalizeFeatures())
z = model.encode(dataset.data.x,dataset.data.edge_index) #encode

emb = TSNE(n_components=2, learning_rate='auto').fit_transform(z.detach().numpy())
labels = dataset.data.y.detach().numpy()
fig, ax = plt.subplots()

number_of_colors = len(np.unique(labels))

color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])
             for i in range(number_of_colors)]
for idx , i in enumerate(np.unique(labels)) :
    emb_ = emb[np.where(labels == i ),:].squeeze()
    ax.scatter(x=emb_[:,0],y=emb_[:,1],c=color[idx], label=i,alpha=0.2)
else :
    ax.legend()
    plt.show()

 

encoding된 결과를 다시 차원 축소를 통해서 본 결과 특정 Node 들끼리는 논문 분류에 맞게 잘 임베딩 되어 있음을 알 수 있습니다.

Conclusion

이번에는 link prediction 구현된 것을 가지고 왔다. 

생각보다 예전에는 어떻게 할 수 있을까 하는 막연한 생각이 있었는데, 생각보다 쉽게 할 수 있어 보인다.

decode의 방식에는 제일 간단한 방법을 한 것 같고, 좀 더 복잡하게도 할 때도 할 수 있게다는 생각이 들었다. 

다만 더 궁금한 것은 실제 데이터에는 더 많은 node와 edge를 가지고 있을 텐데, 이 데이터에 따라서 어떻게 배치로 학습하는지 찾아봤지만 잘 안 나왔고, 새로운 node나 새로운 edge는 계속 생길 텐데, 이것을 어떻게 반영하는지가 아직 미스터리다.

 

 

Reference

 

https://colab.research.google.com/github/AntonioLonga/PytorchGeometricTutorial/blob/main/Tutorial12/Tutorial12%20GAE%20for%20link%20prediction.ipynb#scrollTo=ZEICfIpN4SBS

Tutorial12 GAE for link prediction.ipynb