Preliminaries
PyTorch Geometric official docs
PyTorch Geometric official github repo
Stanford PyTorch Geometric tutorial (YouTube)
몇몇 GNN paper를 차츰 읽어보고, 관련 라이브러리가 있는지 찾아보던 중 알게된 PyTorch Geometric, 일명 PyG.
이때가 아마 22년 9월~10월 경인데, 기억상 PyG가 정식으로 공개된지 얼마 되지 않기도 하였고 한창 개발중이던 때라 왠지 더 눈길이 갔던 것 같다.
그리고 무엇보다 마음에 들었던 점은 기존 PyTorch의 code style과 거의 비슷하여서 익숙해지는데 큰 어려움이 없었던 점이다.
더욱 마음에 들었던 점은 바로 공식 docs에 예제 code와 video tutorial를 제공해주어서, GNN의 기본부터 block building, downstream task, advanced topics까지 필요한 부분을 학습할 수 있는 것이 정말 좋았다.
본 포스팅에서는 PyG를 처음 접했을 때 기본 code block의 구성이 어떻게 되는지, code flow가 어떻게 진행되는지를 간략하게 정리하고자 한다.
Basic Data Handling
PyTorch Geometric official docs에서 제공하는 문서를 보면 PyG에서 데이터를 어떻게 다룰 수 있는지 설명하고 있는데, 간략하게 정리해보면 다음과 같다:
data.x: 각 node마다 가지고 있는 feature를 matrix 형태로 표현. shape는 $(number \ of \ nodes, \ feature \ dim)$가 된다.data.edge_index: 각 node간의 연결 정보를 담고 있는 adj matrix 같은 것으로, 효율성을 위해 COO format으로 표현.data.edge_attr: 각 edge마다 가지고 있는 feature를 마찬가지로 matrix 형태로 표현. shape는 $(number \ of \ edges, \ feature \ dim)$가 된다.data.y: training을 위한 label 정보. shape는 $(number \ of \ nodes, \ *)$ 또는 (graph-level task)에선 $(1, \ *)$가 된다.
단순하다.
실제로 PyG에서 제공하는 Dataset을 받아서 code를 작성하다 보면 대부분 위에 기술한 4가지 인자를 주로 사용하게 된다.
기존 adjacency matrix로는 표현 및 code 활용의 한계가 있어 위와 같은 이중 list 형태의 COO format을 사용하는데, 이게 상당히 직관적이다.
(단순히 첫 번째 row와 두 번째 row 간의 요소들이 서로 연결되어 있다를 의미한다.)
Message Passing
PyG official docs에서도 설명하는 GNN의 핵심인 message passing의 기본적인 골격은 다음과 같다.
$$\mathbf{x}_i^{(k)} = \gamma^{(k)} \left( \mathbf{x}_i^{(k-1)}, \bigoplus_{j \in \mathcal{N}(i)} \, \phi^{(k)}\left(\mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)},\mathbf{e}_{j,i}\right) \right)$$
$\mathbf{x}^{(k-1)}_i \in \mathbb{R}^F$가 node $i$의 feature를 의미하며, $\mathbf{e}_{j,i} \in \mathbb{R}^D$는 node $j$와 $i$간 연결된 edge의 feature를 의미한다.
$\bigoplus$가 바로 각 node와 edge들의 feature를 aggregate하는 함수로, 미분 가능하고 순서 불변(permutation invariant)한 특징을 가진다.
(GraphSAGE paper에서도 언급하였듯이, 주로 활용할 수 있는 함수(연산)는 sum, mean, max 가 있다.)
($\phi$와 $\gamma$도 마찬가지로 MLP와 같은 미분 가능한 함수이다.)
PyG에서 제공하는 torch_geometric.nn의 여러 layer module들은 기본적으로 basic class인 torch_geometric.nn.MessagePassing module을 상속 받으며, MessagePassing은 torch.nn.Module을 상속 받는다.
이 MessagePassing은 기본적으로 4가지의 basic block을 구현해야 하는데, 그 구성은 다음과 같다:
__init__(aggr='mean', 'max', 'add'): feature aggregation의 연산에 있어 어떤 연산을 수행할지를 지정한다.propagate():edge_index와 각 노드들에 전달할 데이터(e.g. feature)를 입력받아 message passing을 위한 초기 단계를 준비한다.message(): target node에 전파할 message를 생성하며,propagate()에 의해 호출된다.update(): 계산된 message를 이용해 embedding을 갱신한다.
즉, __init__(aggr='mean', 'max', 'add')를 통해 $\bigoplus$를 지정하고, propagate(), message()를 통해 $\phi^{(k)}\left(\mathbf{x}_i^{(k-1)}, \mathbf{x}_j^{(k-1)},\mathbf{e}_{j,i}\right) $의 연산을 수행하며 최종적으로 update()를 통해 $\gamma^{(k)}\left( \cdot \right)$의 연산을 수행하는 셈이다.
실제로 PyG github repo에 공개되어 있는 GNN layer들의 구현들을 살펴보면 MessagePassing을 상속받아 상단의 함수들을 구현하여 layer를 구현한다.
Example Code
Stanford PyTorch Geometric tutorial (YouTube)에서 활용한 colab notebook의 code의 일부를 이용해 기본적인 flow를 살펴보고자 한다.
우선 필요한 library를 import 한 후, MessagePassing을 상속받는 CustomConv class를 작성한다.
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch_geometric.nn as pyg_nn
import torch_geometric.utils as pyg_utils
import torch_geometric.transforms as T
from torch_geometric.data import DataLoader
from torch_geometric.data import Planetoid
class CustomConv(pyg_nn.MessagePassing):
def __init__(self, in_channels, out_channels):
# add, mean, max 중 add aggregation으로 지정.
super(CustomConv, self).__init__(aggr='add')
self.lin = nn.Linear(in_channels, out_channels)
self.lin_self = nn.Linear(in_channels, out_channels)
def forward(self, x, edge_index):
# x의 shape은 [N, in_channels]
# edge_index의 shape은 [2, E]
# Add self-loops to the adjacency matrix.
edge_index, _ = pyg_utils.add_self_loops(edge_index)
# Node feature matrix 선형 변환
self_x = self.lin_self(x)
# propagate()을 통해 message passing의 초기 단계 준비
return self_x + self.propagate(edge_index, size=(x.size(0), x.size(0)), x=self.lin(x))
def message(self, x_i, x_j, edge_index, size):
# Message 계산
# x_j의 shape은 [E, out_channels]
# GCN-like message passing
row, col = edge_index
deg = pyg_utils.degree(row, size[0], dtype=x_j.dtype)
deg_inv_sqrt = deg.pow(-0.5)
norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
return x_j
def update(self, aggr_out):
# aggr_out의 shape은 [N, out_channels]
return aggr_out
__init__(aggr='add')를 통해 feature aggregation의 연산 방법을 지정하였다.
기존 PyTorch class들과 마찬가지로 forward()를 통해 주 연산을 수행하고, propagate()를 통해 message()와 update()를 호출하여 node embedding을 계산한다.
PyTorch만으로 작성하는 기본적인 여느 model/layer와 큰 차이가 없고, 단지 사용하는 데이터가 image, word가 아닌 graph가 사용됨에 따라 추가적인 연산(feature aggregation 및 message passing)이 추가된 셈이다.
작성된 CustomConv layer를 활용해 기본적인 model을 작성해보면 다음과 같다:
class GNNStack(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, task='node'):
super(GNNStack, self).__init__()
self.task = task
self.convs = nn.ModuleList()
self.convs.append(self.build_conv_model(input_dim, hidden_dim))
self.lns = nn.ModuleList()
self.lns.append(nn.LayerNorm(hidden_dim))
self.lns.append(nn.LayerNorm(hidden_dim))
for l in range(2):
self.convs.append(self.build_conv_model(hidden_dim, hidden_dim))
# post-message-passing
self.post_mp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim), nn.Dropout(0.25),
nn.Linear(hidden_dim, output_dim))
if not (self.task == 'node' or self.task == 'graph'):
raise RuntimeError('Unknown task.')
self.dropout = 0.25
self.num_layers = 3
def build_conv_model(self, input_dim, hidden_dim):
# refer to pytorch geometric nn module for different implementation of GNNs.
if self.task == 'node':
# 예제 code는 torch_geometric.nn.GCNConv를 사용.
# return pyg_nn.GCNConv(input_dim, hidden_dim)
return CustomConv(input_dim, hidden_dim)
else:
return pyg_nn.GINConv(nn.Sequential(nn.Linear(input_dim, hidden_dim),
nn.ReLU(), nn.Linear(hidden_dim, hidden_dim)))
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
if data.num_node_features == 0:
x = torch.ones(data.num_nodes, 1)
for i in range(self.num_layers):
x = self.convs[i](x, edge_index)
emb = x
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
if not i == self.num_layers - 1:
x = self.lns[i](x)
# graph-level task는 graph-level embedding을 활용하므로,
# global pooling 연산을 이용해 graph-level embedding을 생성한다.
if self.task == 'graph':
x = pyg_nn.global_mean_pool(x, batch)
x = self.post_mp(x)
return emb, F.log_softmax(x, dim=1)
def loss(self, pred, label):
return F.nll_loss(pred, label)
제공된 colab code는 node-level task(node classification, Planetoid dataset)와 graph-level task(graph classification, TU dataset)을 수행하는 예시를 보여준다.
(추후 posting 할 예정이지만) graph-level task는 node-level task와는 다르게 말 그대로 graph 하나 하나를 가지고 downstream task를 수행하는 것으로, node embedding이 아닌 graph embedding이 필요하다.
이에 따라 상기 code에서는 global_mean_pool(x, batch)를 통해 각 graph마다 계산된 node embedding에 대해 global pooling 연산을 수행하여 하나의 embedding, 즉 graph embedding을 생성하였다.
PyTorch로 작성된 여느 model과 마찬가지로 forward()를 통해 주 연산을 수행한다.
앞서 언급하였듯이 단지 사용되는 데이터가 graph이기에 node feature와 연결 정보를 담은 edge_index를 이용해 주 연산을 수행한다.
train 및 test를 수행하는 code도 기존 PyTorch의 code style과 거의 동일하다:
def train(dataset, task):
if task == 'graph':
data_size = len(dataset)
loader = DataLoader(dataset[:int(data_size * 0.8)], batch_size=64, shuffle=True)
test_loader = DataLoader(dataset[int(data_size * 0.8):], batch_size=64, shuffle=True)
else:
test_loader = loader = DataLoader(dataset, batch_size=64, shuffle=True)
# model 선언
model = GNNStack(max(dataset.num_node_features, 1), 32, dataset.num_classes, task=task)
opt = optim.Adam(model.parameters(), lr=0.01)
# train
for epoch in range(200):
total_loss = 0
model.train()
for batch in loader:
print(batch.train_mask, '----')
opt.zero_grad()
embedding, pred = model(batch)
label = batch.y
if task == 'node':
pred = pred[batch.train_mask]
label = label[batch.train_mask]
loss = model.loss(pred, label)
loss.backward()
opt.step()
total_loss += loss.item() * batch.num_graphs
total_loss /= len(loader.dataset)
if epoch % 10 == 0:
test_acc = test(test_loader, model)
print("Epoch {}. Loss: {:.4f}. Test accuracy: {:.4f}".format(
epoch, total_loss, test_acc))
return model
def test(loader, model, is_validation=False):
model.eval()
correct = 0
for data in loader:
with torch.no_grad():
emb, pred = model(data)
pred = pred.argmax(dim=1)
label = data.y
if model.task == 'node':
mask = data.val_mask if is_validation else data.test_mask
# node classification: only evaluate on nodes in test set
pred = pred[mask]
label = data.y[mask]
correct += pred.eq(label).sum().item()
if model.task == 'graph':
total = len(loader.dataset)
else:
total = 0
for data in loader.dataset:
total += torch.sum(data.test_mask).item()
return correct / total
# model train with Cora dataset
dataset = Planetoid(root='/tmp/cora', name='cora')
task = 'node'
model = train(dataset, task)
여느 PyTorch script와 마찬가지로 loader를 통해 dataset을 불러오고, model을 선언하고, optimizer를 선언하고, train 및 test를 진행한다.
학습이 완료된 model을 통해 Cora dataset (citation network)에서의 node classification 결과를 확인해보면 다음과 같은 결과를 얻을 수 있다:
학습된 GNN model을 통해 Cora dataset에서 각 paper(node)가 어느 주제(class)에 속하는지를 분류하였다.
Conclusion
작성하다보니 포스팅이 조금 길어진 듯 하다.
앞선 예시 code를 살펴봐도 알 수 있듯이, PyG는 기존 PyTorch와 거의 완벽하게 호환이 될 정도로 code style이 간편한것이 정말 큰 장점인 듯 하다.
(물론 내부적으로 더 나은 연산을 위해 torch-sparse, torch-scatter와 같은 라이브러리를 별도로 사용하기도 한다.)
전체적인 code style은 기존 PyTorch와 거의 유사하지만, 사용하는 데이터가 graph 데이터 이기에 전 처리 과정 및 model/layer에서 연산 수행을 위한 code style이 조금 다를 뿐이다.
22년, 23년 당시 GNN paper와 공개된 github repo를 잠시 찾아본 적이 있는데, 실제로 새로운 model을 구현한 work 중 PyG를 통해 구현한 경우 위 처럼 MessagePassing을 상속받아 제안한 model/layer를 직접 설계한 모습을 볼 수 있었다.
물론 이렇게 직접 구현할 수 있지만, 라이브러리는 괜히 라이브러리가 아닌 지라 PyG team 에서도 적극적으로 outstanding한 work들의 layer를 직접 구현하여 모듈로 제공하고 있다.
(well-build 된 것이 있다면 역시 그것을 사용하는게...)
그리고 Jure Leskovec 교수님이 강의에서 직접 PyG를 사용하도록 권장(?)하고 있기도 하고, PyG를 개발한 Matthias Fey를 포함한 PyG team들과 함께 kumo라는 벤쳐 기업을 설립하여 그런지,
실제 비즈니스 적 측면에서 PyG를 어떻게 더 효율적으로 최적화하는지에 대한 연구도 함께 진행되고 있다는 점을 23년 Webinar에서 직접 듣기도 한 만큼 DGL과 더불어 유망한 라이브러리로 거듭나고 있는 느낌이 많이 들었다.
만일 GNN을 입문하고자 하는 지인/후배가 있다고 한다면,
나는 물론 실제로 coding이 간편하기도 하고, 무엇보다 양질의 tutorial & video가 정말 잘 제공되는 PyG를 추천할 생각이다.
(그렇다고 DGL은 별로다 라는 것은 아니다 😅 분산 학습과 같이 system 적인 측면에서는 DGL이 아직은 PyG보다 더욱 well-build 되어 있다고 생각하기에...)
E.O.D.
'학업 이야기 > Programming' 카테고리의 다른 글
| PyTorch Geometric(PyG) - Cora EDA & mini-batch testing(with NeighborLoader) (1) | 2025.02.17 |
|---|---|
| PyTorch Geometric(PyG) - NeighborLoader (3) | 2024.08.26 |