torch-geometric

PyTorch Geometric(PyG)

PyG는 PyTorch 위에 구축된 그래프 신경망(Graph Neural Networks)을 위한 표준 라이브러리입니다. 그래프를 위한 데이터 구조, 60개 이상의 GNN 레이어 구현, 확장 가능한 미니배치 학습, 그리고 이종 그래프(heterogeneous graph) 지원을 제공합니다.

설치

torch-geometric 2.7.x (2025년 10월 기준)에서 테스트되었습니다. Python 3.10+ 및 PyTorch 2.6+ 가 필요합니다.

# 1. 먼저 PyTorch를 설치하세요 (사용자의 CUDA/CPU 설정에 맞게 설치 — https://pytorch.org/get-started/locally/ 참조)
uv pip install torch

# 2. 코어 PyG (기본 사용을 위해 확장 휠은 필요하지 않음)
uv pip install torch_geometric

선택적 가속 연산(pyg-lib, torch-scatter, torch-sparse, torch-cluster)은 PyG 2.3부터 기본 사용에 필수가 아닙니다. PyTorch 및 CUDA 버전을 확인한 후 PyG 휠 인덱스에서 버전에 맞는 휠을 설치하세요:

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# 그런 다음 torch+CUDA 조합에 맞는 휠을 설치합니다. 예:
uv pip install pyg-lib torch-scatter torch-sparse torch-cluster \
  -f https://data.pyg.org/whl/torch-2.8.0+cu128.html

버전 확인:

import torch_geometric
print(torch_geometric.__version__)

Conda: pyg conda 채널은 PyTorch >2.5에 대해 더 이상 유지 관리되지 않습니다. 위의 휠 인덱스와 함께 uv pip install을 사용하세요.

PyG 2.7 참고 사항

PyG 2.7은 Python 3.9 및 PyTorch ≤2.5 지원을 중단했습니다. PyTorch 2.6–2.8 호환성 표는 2.7.0 릴리스 노트를 참조하세요. torch_geometric.distributed는 더 이상 사용되지 않습니다. 표준 torch.distributed DDP를 사용하세요 (references/scaling.md 참조).

핵심 개념

그래프 데이터: Data 및 HeteroData

그래프는 Data 객체에 저장됩니다. 주요 속성은 다음과 같습니다:

from torch_geometric.data import Data

data = Data(
    x=node_features,          # [num_nodes, num_node_features]
    edge_index=edge_index,     # [2, num_edges] — COO 형식, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]
    y=labels,                  # 노드 수준 [num_nodes, *] 또는 그래프 수준 [1, *]
    pos=positions,             # [num_nodes, num_dimensions] (포인트 클라우드/공간 데이터용)
)

edge_index 형식은 매우 중요합니다: [2, num_edges] 텐서이며 edge_index[0] = 소스 노드, edge_index[1] = 타겟 노드입니다. 튜플 목록이 아닙니다. 엣지 쌍이 행 형태로 주어지면 전치하고 .contiguous()를 호출하세요:

# 엣지가 [[src1, dst1], [src2, dst2], ...] 형태이면 먼저 전치합니다:
edge_index = edge_pairs.t().contiguous()

무방향 그래프의 경우, 두 방향을 모두 포함해야 합니다: 엣지 (0,1)은 edge_index에 [0,1]과 [1,0]이 모두 필요합니다.

이종 그래프의 경우 HeteroData를 사용하세요 — 아래 '이종 그래프' 섹션을 참조하세요.

데이터셋

PyG는 자동으로 다운로드 및 전처리되는 많은 표준 데이터셋을 제공합니다:

from torch_geometric.datasets import Planetoid, TUDataset

# 단일 그래프 노드 분류 (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]  # train/val/test 마스크가 포함된 단일 그래프

# 다중 그래프 분류 (ENZYMES, MUTAG, IMDB-BINARY 등)
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... 은 개별 그래프입니다

작업별 일반적인 데이터셋:

• 노드 분류: Planetoid (Cora/Citeseer/Pubmed), OGB (ogbn-arxiv, ogbn-products, ogbn-mag)
• 그래프 분류: TUDataset (MUTAG, ENZYMES, PROTEINS, IMDB-BINARY), OGB (ogbg-molhiv)
• 링크 예측: OGB (ogbl-collab, ogbl-citation2)
• 분자: QM7, QM9, MoleculeNet
• 포인트 클라우드/메시: ShapeNet, ModelNet10/40, FAUST

변환(Transforms)

변환은 torchvision 변환과 유사하게 그래프 데이터를 전처리하거나 증강합니다:

import torch_geometric.transforms as T

# 일반적인 변환
T.NormalizeFeatures()    # 노드 특징을 행 단위로 정규화하여 합이 1이 되도록 함
T.ToUndirected()         # 그래프를 무방향으로 만들기 위해 역방향 엣지 추가
T.AddSelfLoops()         # 자기 루프 엣지 추가
T.KNNGraph(k=6)          # 포인트 클라우드 위치에서 k-NN 그래프 구축
T.RandomJitter(0.01)     # 위치에 무작위 노이즈 증강
T.Compose([...])         # 여러 변환 연결

# pre_transform (한 번 적용, 디스크에 저장) 또는 transform (매 접근 시)으로 적용
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

GNN 모델 구축

빠른 시작: 내장 레이어 사용

GNN 구축의 가장 빠른 방법 — torch_geometric.nn에서 conv 레이어를 쌓으세요:

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

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

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

중요: PyG conv 레이어는 활성화 함수를 포함하지 않습니다. 각 레이어 이후 직접 적용해야 합니다. 이는 유연성을 위한 의도된 설계입니다.

Conv 레이어 선택

작업 및 그래프 구조에 따라 선택하세요:

레이어	최적 대상	핵심 아이디어
GCNConv	동종, 준지도 노드 분류	스펙트럴 기반, 차수 정규화 집계
GATConv / GATv2Conv	이웃 중요도가 다양한 경우	어텐션 가중치 메시지
SAGEConv	대규모 그래프, 귀납적 설정	샘플링 친화적, 학습 가능한 집계
GINConv	그래프 분류, 표현력 극대화	WL 테스트만큼 강력
TransformerConv	풍부한 엣지 특징, 복잡한 상호작용	엣지 특징을 포함한 멀티헤드 어텐션
EdgeConv	포인트 클라우드, 동적 그래프	엣지 특징에 대한 MLP (x_i, x_j - x_i)
RGCNConv	다양한 관계 유형의 이종 그래프	관계별 가중치 행렬
HGTConv	이종 그래프	유형별 어텐션

모든 conv 레이어는 최소한 (x, edge_index)를 입력으로 받습니다. 대부분은 엣지 특징을 위한 edge_attr도 받습니다.

지연 초기화(Lazy Initialization)

입력 채널에 -1을 사용하면 PyG가 차원을 자동으로 추론합니다. 특히 이종 모델에서 유용합니다:

conv = SAGEConv((-1, -1), 64)  # 첫 번째 순방향 패스에서 입력 차원이 추론됨
# 지연 모듈 초기화:
with torch.no_grad():
    out = model(data.x, data.edge_index)

고수준 모델 API

일반적인 아키텍처의 경우 PyG는 미리 만들어진 모델 클래스를 제공합니다:

from torch_geometric.nn import GraphSAGE, GCN, GAT, GIN

model = GraphSAGE(
    in_channels=dataset.num_features,
    hidden_channels=64,
    out_channels=dataset.num_classes,
    num_layers=2,
)

MessagePassing을 사용한 사용자 정의 레이어

새로운 GNN 레이어를 구현하려면 MessagePassing을 서브클래싱하세요. 프레임워크는 다음과 같습니다:

1. propagate()는 메시지 전달을 조정합니다.
2. message()는 각 엣지를 따라 어떤 정보가 흐를지 정의합니다 (phi 함수).
3. aggregate()는 각 노드에서 메시지를 결합합니다 (sum/mean/max).
4. update()는 집계된 결과를 변환합니다 (gamma 함수).

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class MyConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # "add", "mean", 또는 "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # 메시지 전달 전 전처리
        x = self.lin(x)
        # 메시지 전달 시작
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: 각 엣지의 소스 노드 특징 [num_edges, features]
        # _j 접미사는 소스 노드를 자동 인덱싱하고, _i는 타겟 노드를 인덱싱함
        return x_j

_i / _j 규칙: propagate()에 전달된 모든 텐서는 message() 시그니처에서 _i (타겟/중심 노드) 또는 _j (소스/이웃 노드)를 추가하여 자동 인덱싱할 수 있습니다. 따라서 propagate에 x=...를 전달하면 message()에서 x_i와 x_j에 접근할 수 있습니다.

전체 GCN 및 EdgeConv 구현 예제는 references/message_passing.md를 참조하세요.

작업별 패턴

노드 분류

# 단일 그래프 전체 배치 학습 (예: Cora)
model.train()
for epoch in range(200):
    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()

# 평가 — train(False)는 모델을 추론 모드로 전환합니다 (드롭아웃/BN 비활성화)
model.train(False)
pred = model(data.x, data.edge_index).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()

그래프 분류

여러 그래프 — DataLoader를 사용한 미니배칭과 그래프 수준 표현을 위한 전역 풀링:

from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool

loader = DataLoader(dataset, batch_size=32, shuffle=True)

class GraphClassifier(torch.nn.Module):
    def __init__(self, in_ch, hidden_ch, out_ch):
        super().__init__()
        self.conv1 = GCNConv(in_ch, hidden_ch)
        self.conv2 = GCNConv(hidden_ch, hidden_ch)
        self.lin = torch.nn.Linear(hidden_ch, out_ch)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        x = global_mean_pool(x, batch)  # [num_graphs_in_batch, hidden_ch]
        return self.lin(x)

# 학습 루프
for data in loader:
    out = model(data.x, data.edge_index, data.batch)
    loss = F.cross_entropy(out, data.y)

PyG의 DataLoader는 여러 그래프를 블록 대각 인접 행렬로 배칭합니다. batch 텐서는 각 노드를 해당 그래프 인덱스에 매핑합니다. 풀링 연산(global_mean_pool, global_max_pool, global_add_pool)은 이를 사용하여 그래프별 집계를 수행합니다.

링크 예측

엣지를 train/val/test로 분할하고, 음성 샘플링 사용:

from torch_geometric.transforms import RandomLinkSplit

transform = RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    is_undirected=True,
    add_negative_train_samples=False,
)
train_data, val_data, test_data = transform(data)

# 노드 인코딩 후 엣지 점수 계산
z = model.encode(train_data.x, train_data.edge_index)
# 양성 엣지
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)

전체 링크 예측 가이드는 references/link_prediction.md를 참조하세요: GAE/VGAE 오토인코더, 전체 학습 루프, 대규모 그래프를 위한 LinkNeighborLoader, 이종 그래프 링크 예측, 평가 지표.

대규모 그래프로 확장

GPU 메모리에 맞지 않는 그래프의 경우 NeighborLoader를 통해 이웃 샘플링을 사용하세요:

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[15, 10],     # 1차 홉에서 15개, 2차 홉에서 10개 이웃 샘플링
    batch_size=128,              # 배치당 시드 노드 수
    input_nodes=data.train_mask, # 샘플링할 노드 지정
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # 손실 계산에는 처음 batch_size개 노드만 사용 (시드 노드)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

NeighborLoader 핵심 사항:

• num_neighbors 리스트 길이는 GNN 깊이(메시지 전달 레이어 수)와 일치해야 합니다.
• 시드 노드는 항상 출력에서 처음 batch.batch_size개 노드입니다.
• batch.n_id는 다시 레이블링된 인덱스를 원래 노드 ID로 매핑합니다.
• Data와 HeteroData 모두에서 작동합니다.
• 링크 예측의 경우 LinkNeighborLoader를 대신 사용하세요.
• 일반적으로 2-3 홉 이상의 샘플링은 불가능합니다 (지수적 폭발).

기타 확장 옵션: ClusterLoader (ClusterGCN), GraphSAINTSampler, ShaDowKHopSampler. 다중 GPU 학습, DDP, PyTorch Lightning 통합, torch.compile 지원은 references/scaling.md를 참조하세요.

이종 그래프

여러 노드 및 엣지 유형이 있는 그래프 (소셜 네트워크, 지식 그래프, 추천 시스템):

from torch_geometric.data import HeteroData

data = HeteroData()

# 노드 특징 — 노드 유형 문자열로 인덱싱
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)

# 엣지 인덱스 — (src_type, edge_type, dst_type) 삼중항으로 인덱싱
data['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))

# 편의 딕셔너리 접근
data.x_dict        # {'user': tensor, 'movie': tensor}
data.edge_index_dict  # {('user','rates','movie'): tensor, ...}
data.metadata()    # ([node_types], [edge_types])

이종 GNN 구축의 세 가지 방법

1. to_hetero()로 자동 변환 — 동종 모델을 작성하고 자동으로 변환:

from torch_geometric.nn import SAGEConv, to_hetero

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

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

model = GNN(64, dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

# 이제 딕셔너리를 입력으로 받습니다:
out = model(data.x_dict, data.edge_index_dict)

이분 그래프 입력 채널에는 (-1, -1)을 사용하세요 (소스와 타겟이 다를 수 있음). 지연 초기화가 나머지를 처리합니다.

2. HeteroConv 래퍼 — 엣지 유형별로 다른 conv:

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv

conv = HeteroConv({
    ('paper', 'cites', 'paper'): GCNConv(-1, 64),
    ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
    ('paper', 'written_by', 'author'): GATConv((-1, -1), 64),
}, aggr='sum')

# 순방향 패스는 딕셔너리를 입력받고 딕셔너리를 반환
# x_dict: {node_type: tensor}, edge_index_dict: {edge_type: tensor}
out_dict = conv(x_dict, edge_index_dict)  # {node_type: tensor}

3. 수동 HeteroData 처리 — 각 노드 유형에 대해 별도 nn.ModuleDict 및 ParameterDict를 사용하거나 to_hetero()와 HeteroConv를 결합하여 사용자 정의 모듈을 만드세요.

이종 그래프 샘플링

HeteroData에 NeighborLoader를 사용하고 각 엣지 유형에 대한 num_neighbors 딕셔너리를 전달하세요:

loader = NeighborLoader(
    data,
    num_neighbors={('user', 'rates', 'movie'): [10],
                   ('user', 'follows', 'user'): [5]},
    batch_size=128,
    input_nodes=('user', data['user'].train_mask),
)

for batch in loader:
    out = model(batch.x_dict, batch.edge_index_dict)
    loss = F.cross_entropy(out[:batch['user'].batch_size],
                           batch['user'].y[:batch['user'].batch_size])

참조

• references/message_passing.md: 사용자 정의 레이어, EdgeConv 예제, 전체 GCN 구현
• references/link_prediction.md: GAE/VGAE, 학습 루프, LinkNeighborLoader, 평가
• references/scaling.md: NeighborLoader, ClusterLoader, GraphSAINT, 멀티 GPU, DDP, torch.compile
• PyG 공식 문서
• PyG 예제 저장소
• PyG GitHub 리포지토리
• PyG 휠 인덱스 — 버전별 호환 가능한 확장 패키지