Graph Convolution Network 구현 wtih DGL

Contents

 

 

Graph Convolution Neural Network(GCN)는 2017년에 처음으로 제안된 모델이다[논문]. 우수한 성능을 보이고 있으며, GNN에 CNN을 접목하기 시작한 계기가 된다. 본 글에서는 이론적인 내용은 자세히 다루지 않고, DGL을 사용해서 GCN을 만들어볼 것이다. GCN에 대한 상세한 내용은 여기를 참고하면 된다.

 

1. Load Package

GCN을 사용하기 위해서는 Pytorch, torch_geometric, Tensorflow, PyG 등 다양한 프레임워크들이 존재하지만 본 글에서는 가장 간단한 DGL을 이용해서 GCN을 구현한다. 구현하기에 앞서 먼저 필요한 패키지들을 호출하자.

 

import numpy as np 
import dgl 
 from dgl.nn import GraphConv 
 
import torch
import torch.nn as nn 
import torch.optim as optim 
import torch.nn.functional as F 

import dgl.data

 

 

2. Load Dataset

원래는 그래프 데이터를 직접 구축해야 되지만, 일단은 DGL에서 제공하는 Cora 데이터를 이용해서 사용하자. Cora 데이터셋은 dgl.data에서 쉽게 호출해서 사용할 수 있다. 

 

dataset = dgl.data.CoraGraphDataset()
print('Number of Categories:', dataset.num_classes)

#  NumNodes: 2708
#  NumEdges: 10556
#  NumFeats: 1433
#  NumClasses: 7
#  NumTrainingSamples: 140
#  NumValidationSamples: 500
#  NumTestSamples: 1000
# Done loading data from cached files.

# Number of Categories: 7


g = dataset[0]

print('Number of nodes:', g.num_nodes())
print('Number of edges:', g.num_edges())

# Number of nodes: 2708
# Number of edges: 10556

 

DGL에서 호출한 데이터는 node data와 edge data로 구성되어 있다. 또, node data는 train mask, val mask, test mask, feats으로 구성되어 있다. mask 데이터는 boolean 값으로 되어 있으며, 만약 node가 train 데이터인 경우 1, 아닌 경우 0으로 표기한다. feats은 node features를 의미한다. 

 

print('Node feature names:', g.ndata.keys())
print('Edge feature names:', g.edata.keys())

# Node feature names: dict_keys(['feat', 'label', 'test_mask', 'train_mask', 'val_mask'])
# Edge feature names: dict_keys(['__orig__'])


print('Number of training nodes:', g.ndata['train_mask'].int().sum().item())
print('Number of validating nodes:', g.ndata['val_mask'].int().sum().item())
print('Number of testing nodes:', g.ndata['test_mask'].int().sum().item())

print('Number of classes:', (g.ndata['label'].max()+1).item())
print('Node feature shape:', g.ndata['feat'].shape)

# Number of training nodes: 140
# Number of validating nodes: 500
# Number of testing nodes: 1000
# Number of classes: 7
# Node feature shape: torch.Size([2708, 1433])

 

 

3. Make Model (GCN)

DGL에서는 매우 쉽게 모델을 만들 수 있도록 하는 프레임워크이다. 그냥 Convolution Layer를 호출하듯 불러와서 사용하면 끝이다.  GCN의 default는 renormalization trick을 사용하는 것이기 때문에 논문을 읽어보고 조금씩 변경해가면서 확인하면 논문과 동일한 모델을 구축할 수 있을 것이다.

class GCN(nn.Module):
    def __init__(self, in_feats, n_hidden, n_classes):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, n_hidden)
        self.conv2 = GraphConv(n_hidden, n_classes)
        self.relu = nn.ReLU()
    
    def forward(self, g, in_feat):
        output = self.conv1(g, in_feat)
        output = self.relu(output)
        output = self.conv2(g, output)
        return output 

models = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes)

 

 

4. Define Function

딥러닝을 학습하기 위해서는 기본적으로 목적 함수, 최적화 함수 등을 정의하여야 한다. 본 글에서는 목적함수를 criterion으로 정의하고 모델의 정확도를 평가하는 함수를 accuracy, 마지막으로 모델을 학습하는 함수를 trainer로 지정하였다. 

def criterion(pred_y, true_y, mask):
    pred_y = pred_y[mask]
    true_y = true_y[mask]
    return F.cross_entropy(pred_y, true_y)

def accuracy(pred_y, true_y, mask):
    pred_y = pred_y[mask]
    true_y = true_y[mask]
    return (pred_y == true_y).float().mean()

def trainer(g, model, n_epoch, device):
    optimizer = optim.Adam(model.parameters())

    best_val_acc = 0 
    best_test_acc = 0 
    g = g.to(device)
    features = g.ndata['feat'].to(device)
    labels = g.ndata['label'].to(device)
    train_mask = g.ndata['train_mask']
    val_mask = g.ndata['val_mask']
    test_mask = g.ndata['test_mask']

    for epoch in range(1, n_epoch+1):
        pred_y = model(g, features).to(device)
        loss = criterion(pred_y, labels, train_mask)
        pred_y = pred_y.argmax(1)
        train_acc = accuracy(pred_y, labels, train_mask)
        val_acc = accuracy(pred_y, labels, val_mask)
        test_acc = accuracy(pred_y, labels, test_mask)

        if best_val_acc < val_acc : 
            best_val_acc = val_acc 
            best_test_acc = test_acc 
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch % 50 == 0 :
            print(f'In epoch [{epoch}/{n_epoch}]\t loss: {loss:.3f}\t train_acc: {train_acc*100:.2f}%\t val_acc: {val_acc*100:.2f}%\t test_acc: {test_acc*100:.2f}%')
            print(f'best val acc: {best_val_acc*100:.2f}%\t best test acc: {best_test_acc*100:.2f}%\n')

 

 

5. Model Training

모델을 학습하는 단계는 위에서 정의한 함수를 단순히 호출해서 사용하면 끝이다. 

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model = GCN(g.ndata['feat'].shape[1], 16, dataset.num_classes).to(device)

trainer(g, model, 1000, device)

# In epoch [700/1000]	 loss: 0.223	 train_acc: 100.00%	 val_acc: 74.20%	 test_acc: 75.70%
# best val acc: 74.40%	 best test acc: 74.60%

# In epoch [750/1000]	 loss: 0.189	 train_acc: 100.00%	 val_acc: 74.60%	 test_acc: 76.20%
# best val acc: 74.60%	 best test acc: 76.20%

# In epoch [800/1000]	 loss: 0.162	 train_acc: 100.00%	 val_acc: 75.00%	 test_acc: 76.00%
# best val acc: 75.00%	 best test acc: 76.10%

# In epoch [850/1000]	 loss: 0.139	 train_acc: 100.00%	 val_acc: 75.20%	 test_acc: 76.00%
# best val acc: 75.20%	 best test acc: 76.10%

# In epoch [900/1000]	 loss: 0.120	 train_acc: 100.00%	 val_acc: 75.40%	 test_acc: 76.00%
# best val acc: 75.60%	 best test acc: 76.00%

# In epoch [950/1000]	 loss: 0.105	 train_acc: 100.00%	 val_acc: 75.40%	 test_acc: 76.10%
# best val acc: 75.60%	 best test acc: 76.00%

# In epoch [1000/1000]	 loss: 0.092	 train_acc: 100.00%	 val_acc: 75.80%	 test_acc: 75.90%
# best val acc: 75.80%	 best test acc: 75.90%

만약 gpu를 사용하고 싶은 경우에는 device를 지정해서 위와 같은 형태로 사용하면 된다. gpu를 사용하지 않거나 없는 경우에는 위 코드를 실행할 때 cpu로 지정되기 때문에 별다른 수정 작업을 수행하지 않아도 된다. 

 

모델의 성능은 75.90%로 막 좋다고는 할 수 없는 성능이 도출되었다. 해당 논문에서의 Cora 데이터셋에 대한 실험 결과는 81%정도로 도출되었으나, 결과의 차이가 조금 존재하는 것 같다. 본 글에서는 initialization과 dropout 등 다양한 튜닝 작업을 수행하지 않았다. 추후 본인이 연습할 때 initialization이나, 목적함수의 learning rate를 조절하는 등의 작업을 수행하면 논문과 동일한 성능이 도출될 것으로 예상된다.