torch-geometric

PyTorch Geometric (PyG)

PyG é a biblioteca padrão para Redes Neurais de Grafos construída sobre o PyTorch. Ela fornece estruturas de dados para grafos, mais de 60 implementações de camadas GNN, treinamento em mini-lotes escalável e suporte a grafos heterogêneos.

Instalação

Testado contra torch-geometric 2.7.x (Out 2025). Requer Python 3.10+ e PyTorch 2.6+.

# 1. Instale o PyTorch primeiro (corresponda à sua configuração CUDA/CPU — veja https://pytorch.org/get-started/locally/)
uv pip install torch

# 2. PyG principal (nenhuma wheel de extensão necessária para uso básico)
uv pip install torch_geometric

Operações aceleradas opcionais (pyg-lib, torch-scatter, torch-sparse, torch-cluster) não são necessárias para uso básico do PyG (desde PyG 2.3). Instale wheels compatíveis com a versão do índice de wheels do PyG após verificar suas versões do PyTorch e CUDA:

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# Em seguida, instale wheels para sua combinação torch+CUDA, por exemplo:
uv pip install pyg-lib torch-scatter torch-sparse torch-cluster \
  -f https://data.pyg.org/whl/torch-2.8.0+cu128.html

Verifique sua versão:

import torch_geometric
print(torch_geometric.__version__)

Conda: o canal conda pyg não é mais mantido para PyTorch >2.5 — use uv pip install e o índice de wheels acima.

Notas do PyG 2.7

PyG 2.7 removeu o suporte para Python 3.9 e PyTorch ≤2.5. Veja as notas de lançamento da versão 2.7.0 para tabelas de compatibilidade com PyTorch 2.6–2.8. torch_geometric.distributed está obsoleto — use DDP torch.distributed padrão (veja references/scaling.md).

Conceitos Centrais

Dados de Grafo: Data e HeteroData

Um grafo reside em um objeto Data. Os atributos principais:

from torch_geometric.data import Data

data = Data(
    x=node_features,          # [num_nodes, num_node_features]
    edge_index=edge_index,     # [2, num_edges] — formato COO, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]
    y=labels,                  # nível de nó [num_nodes, *] ou nível de grafo [1, *]
    pos=positions,             # [num_nodes, num_dimensions] (para nuvens de pontos/espacial)
)

O formato de edge_index é crítico: é um tensor [2, num_edges] onde edge_index[0] = nós de origem, edge_index[1] = nós de destino. NÃO é uma lista de tuplas. Se você tiver pares de arestas como linhas, transpose e chame .contiguous():

# Se as arestas são [[src1, dst1], [src2, dst2], ...] — transpose primeiro:
edge_index = edge_pairs.t().contiguous()

Para grafos não direcionados, inclua ambas as direções: a aresta (0,1) precisa de [0,1] e [1,0] em edge_index.

Para grafos heterogêneos, use HeteroData — veja a seção Grafos Heterogêneos abaixo.

Conjuntos de Dados

PyG inclui muitos conjuntos de dados padrão que baixam e pré-processam automaticamente:

from torch_geometric.datasets import Planetoid, TUDataset

# Classificação de nós em grafo único (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]  # grafo único com máscaras treino/val/teste

# Classificação de múltiplos grafos (ENZYMES, MUTAG, IMDB-BINARY, etc.)
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... são grafos individuais

Conjuntos de dados comuns por tarefa:

• Classificação de nós: Planetoid (Cora/Citeseer/Pubmed), OGB (ogbn-arxiv, ogbn-products, ogbn-mag)
• Classificação de grafos: TUDataset (MUTAG, ENZYMES, PROTEINS, IMDB-BINARY), OGB (ogbg-molhiv)
• Predição de links: OGB (ogbl-collab, ogbl-citation2)
• Molecular: QM7, QM9, MoleculeNet
• Nuvem de pontos/malha: ShapeNet, ModelNet10/40, FAUST

Transformações

Transformações pré-processam ou aumentam dados de grafo, análogas às transformações do torchvision:

import torch_geometric.transforms as T

# Transformações comuns
T.NormalizeFeatures()    # Normaliza linhas das features dos nós para somar 1
T.ToUndirected()         # Adiciona arestas reversas para tornar o grafo não direcionado
T.AddSelfLoops()         # Adiciona arestas de auto-conexão
T.KNNGraph(k=6)          # Constrói grafo k-NN a partir de posições de nuvem de pontos
T.RandomJitter(0.01)     # Aumento de ruído aleatório nas posições
T.Compose([...])         # Encadeia múltiplas transformações

# Aplique como pre_transform (uma vez, salvo em disco) ou transform (a cada acesso)
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

Construindo Modelos GNN

Início Rápido: Usando Camadas Integradas

A maneira mais rápida de construir uma GNN — empilhe camadas conv de torch_geometric.nn:

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

Importante: as camadas conv do PyG NÃO incluem funções de ativação — aplique-as você mesmo após cada camada. Isso é proposital para flexibilidade.

Escolhendo uma Camada Conv

Escolha com base na sua tarefa e estrutura do grafo:

Camada	Melhor para	Ideia chave
GCNConv	Classificação de nós homogênea e semi-supervisionada	Agregação normalizada por grau, inspirada espectralmente
GATConv / GATv2Conv	Quando a importância dos vizinhos varia	Mensagens ponderadas por atenção
SAGEConv	Grafos grandes, configurações indutivas	Amostragem amigável, agregação aprendível
GINConv	Classificação de grafos, maximizando expressividade	Tão poderoso quanto o teste WL
TransformerConv	Features de aresta ricas, interações complexas	Atenção multi-cabeça com features de aresta
EdgeConv	Nuvens de pontos, grafos dinâmicos	MLP em features de aresta (x_i, x_j - x_i)
RGCNConv	Grafos heterogêneos com muitos tipos de relação	Matrizes de peso específicas para relação
HGTConv	Grafos heterogêneos	Atenção específica para tipo

Todas as camadas conv aceitam (x, edge_index) no mínimo. Muitas também aceitam edge_attr para features de aresta.

Inicialização Preguiçosa

Use -1 para canais de entrada para permitir que o PyG infira dimensões automaticamente — especialmente útil para modelos heterogêneos:

conv = SAGEConv((-1, -1), 64)  # Dimensões de entrada inferidas na primeira passagem forward
# Inicialize módulos preguiçosos:
with torch.no_grad():
    out = model(data.x, data.edge_index)

APIs de Modelo de Alto Nível

Para arquiteturas comuns, o PyG fornece classes de modelo prontas:

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,
)

Camadas Personalizadas via MessagePassing

Para implementar uma nova camada GNN, crie uma subclasse de MessagePassing. A estrutura é:

1. propagate() orquestra a passagem de mensagens
2. message() define qual informação flui ao longo de cada aresta (a função phi)
3. aggregate() combina mensagens em cada nó (soma/média/máx)
4. update() transforma o resultado agregado (a função gama)

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", ou "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # Pré-processamento antes da passagem de mensagens
        x = self.lin(x)
        # Inicia a passagem de mensagens
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: features dos nós de origem para cada aresta [num_edges, features]
        # O sufixo _j indexa automaticamente nós de origem, _i indexa nós de destino
        return x_j

A convenção _i / _j: qualquer tensor passado para propagate() pode ser indexado automaticamente anexando _i (nó alvo/central) ou _j (nó fonte/vizinho) na assinatura de message(). Portanto, se você passar x=... para propagate, pode acessar x_i e x_j em message().

Leia references/message_passing.md para exemplos completos de implementação de GCN e EdgeConv.

Padrões Específicos de Tarefas

Classificação de Nós

# Treinamento em lote completo em um único grafo (ex.: 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()

# Avaliação — train(False) coloca o modelo em modo de inferência (desativa dropout/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()

Classificação de Grafos

Múltiplos grafos — use DataLoader para mini-lotes e pooling global para obter representações em nível de grafo:

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)

# Loop de treinamento
for data in loader:
    out = model(data.x, data.edge_index, data.batch)
    loss = F.cross_entropy(out, data.y)

O DataLoader do PyG agrupa múltiplos grafos criando matrizes de adjacência bloco-diagonais. O tensor batch mapeia cada nó ao seu índice de grafo. As operações de pooling (global_mean_pool, global_max_pool, global_add_pool) usam isso para agregar por grafo.

Predição de Links

Divida arestas em treino/validação/teste, use amostragem negativa:

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)

# Codifique nós, depois pontue arestas
z = model.encode(train_data.x, train_data.edge_index)
# Arestas positivas
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)

Leia references/link_prediction.md para o guia completo de predição de links: autoencoders GAE/VGAE, loops de treinamento completos, LinkNeighborLoader para grafos grandes, predição de links heterogênea e métricas de avaliação.

Escalando para Grafos Grandes

Para grafos que não cabem na memória GPU, use amostragem de vizinhos via NeighborLoader:

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[15, 10],     # Amostre 15 vizinhos no salto 1, 10 no salto 2
    batch_size=128,              # Número de nós semente por lote
    input_nodes=data.train_mask, # Quais nós amostrar
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # Use apenas os primeiros batch_size nós para a perda (estes são os nós semente)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

Pontos-chave sobre NeighborLoader:

• O comprimento da lista num_neighbors deve corresponder à profundidade da GNN (número de camadas de passagem de mensagens)
• Os nós semente são sempre os primeiros batch.batch_size nós na saída
• batch.n_id mapeia índices re-rotulados de volta para IDs de nós originais
• Funciona tanto para Data quanto para HeteroData
• Para predição de links, use LinkNeighborLoader
• Amostrar mais de 2-3 saltos é geralmente inviável (explosão exponencial)

Outras opções de escalabilidade: ClusterLoader (ClusterGCN), GraphSAINTSampler, ShaDowKHopSampler. Para treinamento multi-GPU, DDP, integração com PyTorch Lightning e suporte a torch.compile, leia references/scaling.md.

Grafos Heterogêneos

Para grafos com múltiplos tipos de nós e arestas (redes sociais, grafos de conhecimento, recomendação):

from torch_geometric.data import HeteroData

data = HeteroData()

# Features dos nós — indexadas pelo tipo de nó (string)
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)

# Índices de arestas — indexados pelo trio (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))

# Dicionários de conveniência para acesso
data.x_dict        # {'user': tensor, 'movie': tensor}
data.edge_index_dict  # {('user','rates','movie'): tensor, ...}
data.metadata()    # ([node_types], [edge_types])

Três maneiras de construir GNNs heterogêneas

1. Conversão automática com to_hetero() — escreva um modelo homogêneo, converta automaticamente:

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')

# Agora aceita dicionários:
out = model(data.x_dict, data.edge_index_dict)

Use (-1, -1) para canais de entrada bipartidos (origem e destino podem diferir). A inicialização preguiçosa cuida do resto.

2. Wrapper HeteroConv — conv diferente por tipo de aresta:

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

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

3. HGTConv — atenção específica para tipo, incorporada:

from torch_geometric.nn import HGTConv

class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, metadata):
        super().__init__()
        self.lin_dict = torch.nn.ModuleDict()
        for node_type in metadata[0]:
            self.lin_dict[node_type] = torch.nn.Linear(-1, hidden_channels)
        self.conv = HGTConv(hidden_channels, hidden_channels, metadata, num_heads)
        self.lin = torch.nn.Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        x_dict = {
            node_type: self.lin_dict[node_type](x).relu()
            for node_type, x in x_dict.items()
        }
        x_dict = self.conv(x_dict, edge_index_dict)
        return {node_type: self.lin(x) for node_type, x in x_dict.items()}

Mini-lotes heterogêneos com amostragem

O HeteroData funciona com NeighborLoader:

loader = NeighborLoader(
    data,
    num_neighbors={key: [15, 10] for key in data.edge_types},
    batch_size=128,
    input_nodes=('user', data['user'].train_mask),
)

Conjuntos de Dados Personalizados

Para conjuntos de dados que não estão incluídos no PyG:

from torch_geometric.data import Data, InMemoryDataset
import os.path as osp

class MyDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        # Carregue de arquivo processado se existir
        self.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['my_data.pt']  # ou [] se gerar programaticamente

    @property
    def processed_file_names(self):
        return ['data.pt']

    def process(self):
        data_list = []
        # Gere/procure dados aqui
        for i in range(100):
            data = Data(x=..., edge_index=..., y=...)
            data_list.append(data)

        if self.pre_transform is not None:
            data_list = [self.pre_transform(d) for d in data_list]

        self.save(data_list, self.processed_paths[0])

# Uso
dataset = MyDataset(root='./data')

Para conjuntos de dados grandes (que não cabem na RAM), crie uma subclasse de Dataset em vez de InMemoryDataset — consulte references/datasets.md.

Depuração e Melhores Práticas

Prevenção de Vazamento de Informação

O vazamento de informação entre treino e teste em grafos é sutil. Diferente de dados tabulares, os nós de teste são observados durante as agregações de mensagens dos nós de treino. Para lidar com isso:

• Classificação de nós transdutiva (treino+teste no mesmo grafo): Use máscaras train_mask/test_mask. Calcule a perda apenas nos nós de treino, mas o grafo inteiro é usado para passagem de mensagens. Isso é intencional no aprendizado transdutivo em grafos.
• Classificação de nós indutiva: Use amostragem de vizinhos (NeighborLoader), que restringe o campo receptivo aos nós amostrados. Apenas os nós semente contribuem para a perda.
• Predição de links: Sempre divida arestas em treino/validação/teste com RandomLinkSplit. Nunca treine com arestas de teste visíveis. Use arestas negativas não observadas para treino.

Erros Comuns e Correções

Sintoma	Causa Provável	Correção
"edge_index must be a long tensor"	dtype incorreto	edge_index = edge_index.long()
Erro de forma: [num_edges, 2] vs [2, num_edges]	edge_index não transposto	edge_pairs.t().contiguous()
Divergência de NaN na perda	Gradientes explodindo	Adicione normalização (LayerNorm, BatchNorm) ou reduza a taxa de aprendizado
O gradiente não flui (loss não diminui)	ReLU/módulos congelados ou nenhuma ativação entre camadas conv	Verifique as funções de ativação
Erro com HeteroData: "key not found"	Chave de tipo de aresta/nó ausente	Verifique data.metadata(); garanta que x_dict e edge_index_dict correspondam
O modelo ignora features de aresta	edge_attr não passado para a camada conv	Algumas camadas precisam de edge_attr explicitamente — verifique a documentação

Principais Melhores Práticas

1. Sempre verifique as formas — print(data), data.x.shape, data.edge_index.shape
2. Use inicialização preguiçosa (-1, -1) para modelos heterogêneos/bipartidos
3. Para treinamento de vários grafos, sempre passe o tensor batch para operações de pooling
4. Avalie em modo de inferência, model.train(False) ou torch.no_grad()
5. Para conjuntos de dados grandes, use amostragem (NeighborLoader, LinkNeighborLoader)
6. Para perda de classificação de nós, máscara com train_mask antes do cross_entropy
7. Para grafos não direcionados, inclua arestas reversas ou use T.ToUndirected()
8. NVlabs/Sionna para grafos de rádio, PyTorch Geometric temporal para GNNs temporais
9. Consulte a documentação do PyG para questões envolvendo conjuntos de dados ou operações específicas