Deep Graph Library 소개
GNN , DGL , Learning on graphs
Introduction
우리가 흔히 알고 있는 인공 신경망에는 가장 기본적인 Fully-connected network 그리고 Convolutional Neural network (CNN)나 Recurrent Neural network (RNN) 등이 있습니다. 이러한 인공 신경망들은 보통 벡터나 행렬 형태로 input이 주어집니다. 하지만 input이 그래프인 경우에는 벡터나 행렬 형태로 나타내는 대신에 Graph Neural Network (GNN)를 사용할 수 있습니다.
GNN의 기본 원리와 간단한 예시에 대해서는 지난 글 1 에서, GNN의 시간 및 공간 복잡도에 대해서는 지난 글 2에서 알아보았습니다. 이번에는 GNN을 조금 더 쉽게 사용할 수 있는 라이브러리인 Deep Graph Library (DGL) 에 대해서 알아보도록 하겠습니다.
Deep Graph Library
DGL은 그래프 형태의 데이터에서 딥러닝을 사용하는데에 특화된 파이썬 패키지입니다. DGL은 아래와 같은 특징을 가집니다:
- NLP 모델을 제공하는 Hugging Face처럼 DGL에서는 이미 구현된 여러 GNN 모델들을 제공하고, node classification / link prediction / graph classification 과 같은 여러 그래프 관련 task에서 사용할 수 있는 데이터셋도 제공합니다.
- 그래프 형태의 데이터에서 사용할 수 있는 또 다른 라이브러리인 PyTorch Geometric과 다르게 PyTorch 뿐만 아니라 Tensorflow나 Apache MXNet 또한 지원한다는 장점이 있습니다. (다만, 앞으로의 설명은 PyTorch 기준으로 설명합니다)
- 그래프 형태의 데이터에서 message passing 등의 필요한 연산을 빠르게 수행할 수 있도록 수행 시간 및 메모리 사용량 측면에서 최적화 되어있습니다.
DGL을 설치하는 방법은 간단합니다. Get started 페이지에 친절하게 PyTorch처럼 사용하는 CUDA 버전 / 패키지 관리자 / 운영체제 / Python 버전만 알려주면 어떤 명령어로 DGL을 설치해야하는지 알려줍니다. 예를 들어, conda를 사용하는 경우
conda install -c dglteam dgl-cuda[CUDA 버전]
와 같이 설치하면 되고, pip을 사용하는 경우에는
pip install dgl-cu[소수점을 제거한 CUDA 버전]
명령어를 이용하면 됩니다. CUDA를 사용하지 않는 버전은 하이픈 뒤를 모두 지우고 dgl로 설치하면 됩니다.
그래프 생성하기
GNN을 사용하기 위해서는 그래프가 필요합니다. DGL에서는 dgl.graph 함수를 호출함으로써 유향 그래프를 만들 수 있습니다. 함수의 인자는 아래와 같습니다:
- data : 그래프의 간선 정보를 ($U$, $V$)와 같은 형태로 넣어주어야 합니다. U와 V는 각각 텐서나 numpy.array나 list와 같이 iterable한 1-d 자료구조를 사용하여야 하며 길이가 같아야 합니다. 모든 i에 대해 $U[i]$에서 $V[i]$로 가는 간선을 생성합니다.
- num_nodes (선택 사항) : 그래프의 노드 개수를 넣어주면 됩니다. 기본 값은 (data에서 가장 큰 index) + 1입니다.
- idtype (선택 사항) : 노드나 간선 아이디를 32비트 정수형으로 할지 64비트 정수형으로 할지 결정하는 인자입니다. int32 또는 int64로 넣어주면 됩니다.
- device (선택 사항) : 그래프 데이터에서 사용하는 텐서들의 device를 결정합니다. 기본 값은 cpu이며, PyTorch에서 device를 설정하는 것처럼 설정하시면 됩니다.
예를 들어, 세 간선 (0, 1), (1, 2), (2, 0)으로 이루어진 그래프는 아래와 같이 만들 수 있습니다.
import dgl
import torch
g = dgl.graph(data=(torch.LongTensor([0, 1, 2]), torch.LongTensor([1, 2, 0])), num_nodes=3, device=device)
기본적으로 dgl.graph는 유향 그래프를 생성하는데, dgl.to_bidirected를 이용하여 무향 그래프로 바꾸어 줄 수 있습니다.
g = dgl.graph(data=(torch.LongTensor([0, 1, 2]), torch.LongTensor([1, 2, 0])), num_nodes=3)
g = dgl.to_bidirected(g).to(device) # equivalent to dgl.graph(data=(torch.LongTensor([0, 1, 2, 1, 2, 0]), torch.LongTensor([1, 2, 0, 0, 1, 2])))
다른 라이브러리에서 만들어진 그래프를 처음부터 다시 만들 필요 없이 dgl의 그래프 형태로 바꾸어주는 기능도 존재합니다. SciPy의 sparse matrix 형태의 데이터는 dgl.from_scipy, NetworkX의 그래프 형태 데이터는 dgl.from_networkx 함수를 이용하여 불러올 수 있습니다.
import networkx as nx
from scipy.sparse import coo_matrix
_g_sp = coo_matrix(([1] * 6, ([0, 0, 1, 1, 2, 2], [1, 2, 0, 2, 0, 1])), shape=(3, 3))
g_sp = dgl.from_scipy(_g_sp)
_g_nx = nx.DiGraph([(0, 1), (0, 2), (1, 2), (1, 0), (2, 0), (2, 1)])
g_nx = dgl.from_networkx(_g_nx)
feature가 있는 그래프를 만들기 위해서는 만들어진 그래프의 노드나 간선에 feature를 할당해주어야 합니다. 모든 노드에 대하여 feature를 넣어주기 위해서는 첫 번째 차원이 그래프의 노드 개수인 텐서가 필요하고, 모든 엣지에 대하여 feature를 넣어주려면 첫 번째 차원이 그래프의 간선 개수인 텐서가 필요합니다. 그래프의 노드 개수와 간선 개수는 각각 g.num_nodes(), g.num_edges() 메서드를 호출하여 구할 수 있습니다. feature는 dictionary 형태로 관리되며, feature를 넣어줄 때와 사용할 때 모두 g.ndata['feature 이름'], g.edata['feature 이름']과 같이 사용하면 됩니다.
g.ndata['x'] = torch.ones(g.num_nodes(), 32, device=device)
g.edata['weights'] = torch.rand(g.num_edges(), device=device)
Message Passing 사용
대부분의 GNN에서 message passing은 반드시 필요합니다. message passing을 통해서 이웃 노드의 feature를 가지고 오고, 가지고 온 feature들을 통해서 그 노드의 새로운 embedding을 결정합니다. DGL에서는 이웃 노드의 feature를 message 형태로 모으는 과정을 message function, 모은 message를 가공하여 새로운 vector를 만들어내는 과정을 reduce function을 사용하여 처리합니다. message function과 reduce function을 동시에 실행시키려면 update_all(message_function, reduce_function) 함수를 사용하면 됩니다.
기본적인 message function이나 reduce function은 dgl.function에 이미 정의되어 있으므로, 별도로 함수를 구현할 필요없이 그대로 가져다 사용하면 됩니다. 대표적으로 많이 사용되는 built-in function들은 다음과 같습니다:
copy_u(h, m): 출발 노드의 ‘h’ feature를 message의 ‘m’ feature에 넣어서 도착 노드에 전달합니다.copy_e(h, m): 간선의 ‘h’ feature를 message의 ‘m’ feature에 넣어서 도착 노드에 전달합니다.u_mul_e(hu, he, m): 출발 노드의 ‘hu’ feature와 간선의 ‘he’ feature를 곱하여 message의 ‘m’ feature에 넣어 도착 노드에 전달합니다.u_add_v(hu, hv, m): 출발 노드의 ‘hu’ feature와 도착 노드의 ‘hv’ feature를 더하여 message의 ‘m’ feature에 넣어 도착 노드에 전달합니다.sum(m, h): 도착 노드가 받은 message의 ‘m’ feature들을 전부 더하여 도착 노드의 ‘h’ feature로 저장합니다.mean(m, h): 도착 노드가 받은 message의 ‘m’ feature들을 전부 평균 내어 도착 노드의 ‘h’ feature로 저장합니다.max(m, h): 도착 노드가 받은 message의 ‘m’ feature들에서 element-wise max 연산을 수행하여 도착 노드의 ‘h’ feature로 저장합니다.min(m, h): 도착 노드가 받은 message의 ‘m’ feature들을 element-wise min 연산을 수행하여 도착 노드의 ‘h’ feature로 저장합니다.
이외에도 다양한 함수들이 존재하며, message function에서 u는 출발 노드, v는 도착 노드를 그리고 e는 두 노드를 잇는 간선을 의미하므로 상황에 따라 알맞은 함수를 사용하면 됩니다.
만약에 원하는 built-in function이 없다면, 함수를 직접 구현하여 사용하는 방법도 있습니다. 기본적으로 custom message function과 reduce function의 형태는 아래와 같이 나타낼 수 있습니다.
def message_func(edges):
return {'m': edges.src['h'] + edges.dst['h']}
def reduce_func(nodes):
return {'h': torch.max(nodes.mailbox['m'], dim=1)}
위의 예시는 각각 built-in function의 u_add_v('h', 'h', 'm'), max('m', 'h')와 동일한 동작을 합니다. message function에서는 출발 노드, 도착 노드, 간선의 정보를 각각 edges.src, edges.dst, edges.data를 통하여 접근할 수 있으며, reduce function에서는 nodes.mailbox를 이용하여 도착 노드에 모인 message들에 접근할 수 있습니다.
Layer 사용하기
DGL은 dgl.nn을 통해서 미리 구현된 모델들을 제공합니다. 예를 들어 GraphSAGE layer를 사용하려고 한다면 GraphSAGE를 따로 구현할 필요 없이 dgl.nn.SAGEConv를 사용하면 되고, Sum pooling을 마지막에 사용하여 그래프의 feature를 구하려고 한다면 dgl.nn.SumPooling을 사용하면 됩니다. 이렇게 이미 만들어진 레이어를 가져다가 사용하는 것은 매우 간단하기 때문에, 여기서는 custom GNN layer를 어떻게 만드는지에 대해서 다루도록 하겠습니다.
Custom GNN의 구조는 아래처럼 나타낼 수 있고, PyTorch에서 custom 모듈을 만드는 방법과 거의 유사합니다. __init__을 통해서 모듈을 초기화하고, forward를 통해서 그래프와 feature를 입력받아서 원하는 대로 처리하면 됩니다.
from torch import nn
class MyConv(nn.Module):
def __init__(self, ...):
# code for initialiation
def forward(self, graph, feats):
# code for running GNN
GraphSAGE layer는 이미 SAGEConv로 구현되어 있지만, GraphSAGE와 max aggregator를 사용하는 모듈을 간략하게 직접 구현해봅시다. GraphSAGE는 각 layer에서 먼저 input feature를 linear layer를 통하여 transform하고, aggregation을 진행합니다. 그 후 max function으로 reduce를 하고 이전에 도착 노드가 가지고 있던 feature와 함께 다시 한 번 transform하여 그 노드의 임베딩을 결정합니다. 이러한 일련의 과정을 아래의 코드처럼 나타낼 수 있습니다.
from torch import nn
import dgl.function as fn
class MyConv(nn.Module):
def __init__(self, in_feats, out_feats, norm, act):
self.linear1 = nn.Linear(in_feats, in_feats, bias=True)
self.linear2 = nn.Linear(in_feats + in_feats, out_feats, bias=True)
def forward(self, graph, feats):
# transform initial feature
graph.ndata['h'] = self.linear1(feats)
# define message function and reduce function then apply
message_func, reduce_func = fn.copy_u('h', 'm'), fn.max('m', 'h')
graph.update_all(message_func, reduce_func)
# concatenate aggregated feature and update the embedding
h = self.linear2(torch.cat((feats, graph.ndata['h']), dim=-1))
return h
Dataset 사용하기
DGL은 기본적으로 많이 사용되는 Node prediction / edge prediction / graph prediction 데이터셋을 제공합니다. 데이터셋의 리스트는 여기서 확인할 수 있습니다. 또한, DGL에서 기본적으로 제공하는 데이터셋 이외에도 Open Graph Benchmark (OGB)에서도 기본적으로 DGL 데이터셋 형태를 지원합니다! 그러므로 OGB에서 사용하는 데이터셋들도 가져다 쓸 수 있으며, 데이터셋을 불러오는 방법은 dataset 링크에 자세히 소개되어 있습니다.
Model 사용하기와 마찬가지로 여기서는 Dataset을 직접 만드는 방법에 대하여 알아보겠습니다. DGL에서는 dgl.data.DGLDataset class를 제공하며, 데이터셋을 불러올 때 일련의 과정을 순차적으로 수행합니다. 먼저 has_cache() 메서드를 통하여 데이터셋이 이미 디스크에 가공된 상태로 저장되어 있는지 확인합니다. 데이터셋이 디스크에 저장되어 있지 않다면, 먼저 download() 메서드에서 url에 있는 데이터를 raw_dir에 다운로드받고, process() 메서드를 통해 데이터를 가공합니다. 그 후 save() 메서드에서 데이터를 디스크의 save_dir에 저장합니다. 디스크에 저장되었다면, load() 메서드에서 디스크로부터 데이터셋을 불러옵니다.
그러므로, 커스텀 데이터셋을 만들기 위해서는 위에서 언급된 다섯 개의 함수의 구현과 함께 PyTorch에서도 데이터셋을 만들 때 필요한 __getitem__(index)과 __len__()이 필요하고, url, raw_dir, 그리고 save_dir을 지정해주어야 합니다. 위의 내용을 포괄하는 기본적인 데이터셋 구조는 여기를 참고해주세요.
has_cache()
디스크의 save_dir에 올바르게 가공된 데이터가 저장되어있는지 확인하면 됩니다. 크기가 작은 데이터셋의 경우에는 결과값을 항상 False만 return하게 하여서 매번 데이터셋을 가공하게 하여도 무방합니다.
download()
데이터가 이미 디스크에 있다면, raw_dir을 데이터가 있는 위치로 지정하면 download를 별도로 구현할 필요가 없습니다. 만약 인터넷을 통해서 다운로드 받아야한다면, dgl.data.utils에 있는 download 함수나 다른 다운로드 관련 파이썬 라이브러리를 통하여 데이터셋을 다운로드 받으면 됩니다.
def download(self):
# 데이터를 self.raw_dir 위치에 다운로드 받습니다
# (option) 필요하다면 유효성 체크를 해줍니다
pass
process()
다운로드 받은 데이터를 가공하는데, 보통 데이터로부터 그래프를 생성하고, train / validation / test split을 나누고, 각 split에 해당하는 데이터들을 학습에 사용할 수 있도록 가공합니다. 어떤 task를 다루느냐에 따라 process 함수의 구현이 달라지게 됩니다. 예를 들어, node classification의 경우에는 하나의 큰 그래프를 저장하고, 노드를 이제 학습에 사용할 수 있도록 알맞게 나누면 되고, graph classification의 경우에는 그래프들과 feature들을 가공한 후 알맞게 데이터셋을 나누게 됩니다.
save() & load()
process에서 가공한 데이터를 디스크에 저장합니다. Pickle을 사용하여 저장할 수도 있으며 dgl에서 제공하는 저장 관련 함수들(dgl.save_graphs, dgl.data.utils.save_info) 을 사용할 수도 있습니다. 데이터 셋을 불러오는 과정은 저장하는 과정의 반대로 수행하면 됩니다. 저장할 때 사용한 방법과 맞게 데이터셋을 불러오면 됩니다.
getitem, len
__getitem__(index)과 __len__()의 경우에는 PyTorch와 동일합니다. __len__()에서는 데이터셋에 포함된 data의 수를, __getitem__(index)는 index에 해당하는 data를 return하게 구현하면 됩니다.
마무리
지금까지 DGL을 사용하여 그래프 생성 방법, Message passing 진행 과정, 모델/데이터셋 사용 및 커스텀 모델/데이터셋을 만드는 방법에 대하여 알아보았습니다. GNN 모델을 학습하는 것은 일반적인 training 과정과 동일하므로 생략하였으며, 이제 대부분의 GNN 모델을 구현하고 사용할 수 있을 것입니다. Heterogenous 그래프에서 training이나 stochastic training 등에 대한 내용에 대해서 알아보고 싶다면, 공식 Documentation 을 참고해주세요.