torch-geometric

PyTorch Geometric (PyG)

PyG es la biblioteca estándar para Redes Neuronales de Grafos construida sobre PyTorch. Proporciona estructuras de datos para grafos, más de 60 implementaciones de capas GNN, entrenamiento mini-batch escalable y soporte para grafos heterogéneos.

Instalación

Probado contra torch-geometric 2.7.x (Oct 2025). Requiere Python 3.10+ y PyTorch 2.6+.

# 1. Instala PyTorch primero (ajusta a tu configuración CUDA/CPU — ver https://pytorch.org/get-started/locally/)
uv pip install torch

# 2. PyG básico (no se requieren extension wheels para uso básico)
uv pip install torch_geometric

Los paquetes opcionales de operaciones aceleradas (pyg-lib, torch-scatter, torch-sparse, torch-cluster) no son necesarios para el uso básico de PyG (desde PyG 2.3). Instala wheels con la versión adecuada desde el índice de wheels de PyG después de verificar tus versiones de PyTorch y CUDA:

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# Luego instala las wheels para tu combinación de torch+CUDA, por ejemplo:
uv pip install pyg-lib torch-scatter torch-sparse torch-cluster \
  -f https://data.pyg.org/whl/torch-2.8.0+cu128.html

Verifica tu versión:

import torch_geometric
print(torch_geometric.__version__)

Conda: el canal conda pyg ya no se mantiene para PyTorch >2.5 — usa uv pip install y el índice de wheels anterior.

Notas de PyG 2.7

PyG 2.7 eliminó el soporte para Python 3.9 y PyTorch ≤2.5. Consulta las notas de la versión 2.7.0 para las tablas de compatibilidad con PyTorch 2.6–2.8. torch_geometric.distributed está obsoleto — usa torch.distributed DDP estándar (ver references/scaling.md).

Conceptos Fundamentales

Datos de Grafo: Data y HeteroData

Un grafo reside en un objeto Data. Los atributos clave:

from torch_geometric.data import Data

data = Data(
    x=caracteristicas_nodos,          # [num_nodos, num_caracteristicas_nodo]
    edge_index=indice_aristas,         # [2, num_aristas] — formato COO, dtype=torch.long
    edge_attr=caracteristicas_aristas, # [num_aristas, num_caracteristicas_arista]
    y=etiquetas,                       # a nivel de nodo [num_nodos, *] o a nivel de grafo [1, *]
    pos=posiciones,                    # [num_nodos, num_dimensiones] (para nubes de puntos/espacial)
)

El formato edge_index es crítico: es un tensor [2, num_aristas] donde edge_index[0] = nodos fuente, edge_index[1] = nodos destino. NO es una lista de tuplas. Si tienes pares de aristas como filas, transpón y llama a .contiguous():

# Si las aristas son [[src1, dst1], [src2, dst2], ...] — transpón primero:
edge_index = pares_aristas.t().contiguous()

Para grafos no dirigidos, incluye ambas direcciones: la arista (0,1) necesita tanto [0,1] como [1,0] en edge_index.

Para grafos heterogéneos, usa HeteroData — consulta la sección de Grafos Heterogéneos a continuación.

Conjuntos de Datos

PyG incluye muchos conjuntos de datos estándar que se descargan y preprocesan automáticamente:

from torch_geometric.datasets import Planetoid, TUDataset

# Clasificación de nodos en un solo grafo (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]  # un solo grafo con máscaras train/val/test

# Clasificación multi-grafo (ENZYMES, MUTAG, IMDB-BINARY, etc.)
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... son grafos individuales

Conjuntos de datos comunes por tarea:

• Clasificación de nodos: Planetoid (Cora/Citeseer/Pubmed), OGB (ogbn-arxiv, ogbn-products, ogbn-mag)
• Clasificación de grafos: TUDataset (MUTAG, ENZYMES, PROTEINS, IMDB-BINARY), OGB (ogbg-molhiv)
• Predicción de enlaces: OGB (ogbl-collab, ogbl-citation2)
• Moleculares: QM7, QM9, MoleculeNet
• Nube de puntos/malla: ShapeNet, ModelNet10/40, FAUST

Transformaciones

Las transformaciones preprocesan o aumentan los datos del grafo, análogas a las transformaciones de torchvision:

import torch_geometric.transforms as T

# Transformaciones comunes
T.NormalizeFeatures()    # Normaliza filas de características de nodos para que sumen 1
T.ToUndirected()         # Añade aristas inversas para hacer el grafo no dirigido
T.AddSelfLoops()         # Añade aristas de autobucle
T.KNNGraph(k=6)          # Construye grafo k-NN a partir de posiciones de nube de puntos
T.RandomJitter(0.01)     # Aumento de ruido aleatorio en posiciones
T.Compose([...])         # Encadena múltiples transformaciones

# Aplica como pre_transform (una vez, guardado en disco) o transform (cada acceso)
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

Construcción de Modelos GNN

Inicio Rápido: Usando Capas Incorporadas

La forma más rápida de construir una GNN — apila capas 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: las capas conv de PyG NO incluyen funciones de activación — aplícalas tú mismo después de cada capa. Esto es intencional para mayor flexibilidad.

Elegir una Capa Conv

Elige según tu tarea y estructura del grafo:

Capa	Mejor para	Idea clave
GCNConv	Grafos homogéneos, clasificación de nodos semisupervisada	Agregación normalizada por grado, inspirada en espectral
GATConv / GATv2Conv	Cuando la importancia de los vecinos varía	Mensajes ponderados por atención
SAGEConv	Grafos grandes, entornos inductivos	Agregación aprendible, amigable con muestreo
GINConv	Clasificación de grafos, maximización de expresividad	Tan potente como la prueba WL
TransformerConv	Características de arista ricas, interacciones complejas	Atención multi-cabeza con características de arista
EdgeConv	Nubes de puntos, grafos dinámicos	MLP sobre características de arista (x_i, x_j - x_i)
RGCNConv	Grafos heterogéneos con muchos tipos de relaciones	Matrices de pesos específicas para cada relación
HGTConv	Grafos heterogéneos	Atención específica por tipo

Todas las capas conv aceptan al menos (x, edge_index). Muchas también aceptan edge_attr para características de arista.

Inicialización Perezosa

Usa -1 para los canales de entrada y deja que PyG infiera las dimensiones automáticamente — especialmente útil para modelos heterogéneos:

conv = SAGEConv((-1, -1), 64)  # Dimensiones de entrada inferidas en el primer forward
# Inicializa módulos perezosos:
with torch.no_grad():
    out = model(data.x, data.edge_index)

APIs de Modelos de Alto Nivel

Para arquitecturas comunes, PyG proporciona clases de modelo listas para usar:

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

Capas Personalizadas mediante MessagePassing

Para implementar una capa GNN novedosa, subclasea MessagePassing. El marco es:

1. propagate() orquesta el paso de mensajes
2. message() define qué información fluye a lo largo de cada arista (la función phi)
3. aggregate() combina los mensajes en cada nodo (suma/media/máximo)
4. update() transforma el resultado agregado (la función gamma)

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

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

    def forward(self, x, edge_index):
        # Preprocesamiento antes del paso de mensajes
        x = self.lin(x)
        # Inicia el paso de mensajes
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: características de los nodos fuente para cada arista [num_aristas, características]
        # El sufijo _j auto-indexa nodos fuente, _i indexa nodos destino
        return x_j

La convención _i / _j: cualquier tensor pasado a propagate() puede ser auto-indexado añadiendo _i (nodo destino/central) o _j (nodo fuente/vecino) en la firma de message(). Así que si pasas x=... a propagate, puedes acceder a x_i y x_j en message().

Lee references/message_passing.md para los ejemplos completos de implementación de GCN y EdgeConv.

Patrones Específicos por Tarea

Clasificación de Nodos

# Entrenamiento con lote completo en un solo grafo (ej., 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()

# Evaluación — train(False) pone el modelo en modo inferencia (desactiva 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()

Clasificación de Grafos

Múltiples grafos — usa DataLoader para mini-lotes y pooling global para obtener representaciones a nivel 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 ClasificadorGrafos(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_grafos_en_lote, hidden_ch]
        return self.lin(x)

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

El DataLoader de PyG agrupa múltiples grafos creando matrices de adyacencia en bloque diagonal. El tensor batch mapea cada nodo a su índice de grafo. Las operaciones de pooling (global_mean_pool, global_max_pool, global_add_pool) lo usan para agregar por grafo.

Predicción de Enlaces

Divide las aristas en entrenamiento/validación/prueba, usa muestreo negativo:

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)

# Codifica nodos, luego puntúa aristas
z = model.encode(train_data.x, train_data.edge_index)
# Aristas positivas
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)

Lee references/link_prediction.md para la guía completa de predicción de enlaces: autoencoders GAE/VGAE, bucles completos de entrenamiento, LinkNeighborLoader para grafos grandes, predicción de enlaces heterogénea y métricas de evaluación.

Escalado a Grafos Grandes

Para grafos que no caben en la memoria de la GPU, usa muestreo de vecinos mediante NeighborLoader:

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[15, 10],     # Muestrea 15 vecinos en el salto 1, 10 en el salto 2
    batch_size=128,              # Número de nodos semilla por lote
    input_nodes=data.train_mask, # De qué nodos muestrear
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # Solo usa los primeros batch_size nodos para la pérdida (estos son los nodos semilla)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

Puntos clave sobre NeighborLoader:

• La longitud de la lista num_neighbors debe coincidir con la profundidad de la GNN (número de capas de paso de mensajes)
• Los nodos semilla son siempre los primeros batch.batch_size nodos en la salida
• batch.n_id mapea los índices reetiquetados de vuelta a los ID de nodo originales
• Funciona tanto para Data como para HeteroData
• Para predicción de enlaces, usa LinkNeighborLoader
• Muestrear más de 2-3 saltos generalmente no es factible (explosión exponencial)

Otras opciones de escalabilidad: ClusterLoader (ClusterGCN), GraphSAINTSampler, ShaDowKHopSampler. Para entrenamiento multi-GPU, DDP, integración con PyTorch Lightning y soporte torch.compile, lee references/scaling.md.

Grafos Heterogéneos

Para grafos con múltiples tipos de nodos y aristas (redes sociales, grafos de conocimiento, recomendación):

from torch_geometric.data import HeteroData

data = HeteroData()

# Características de nodos — indexadas por cadena de tipo de nodo
data['usuario'].x = torch.randn(1000, 64)
data['pelicula'].x = torch.randn(500, 128)

# Índices de aristas — indexados por triplete (tipo_fuente, tipo_arista, tipo_destino)
data['usuario', 'califica', 'pelicula'].edge_index = torch.randint(0, 500, (2, 3000))
data['usuario', 'sigue_a', 'usuario'].edge_index = torch.randint(0, 1000, (2, 5000))

# Diccionarios de acceso convenientes
data.x_dict        # {'usuario': tensor, 'pelicula': tensor}
data.edge_index_dict  # {('usuario','califica','pelicula'): tensor, ...}
data.metadata()    # ([tipos_nodo], [tipos_arista])

Tres formas de construir GNNs heterogéneas

1. Auto-conversión con to_hetero() — escribe un modelo homogéneo, conviértelo automáticamente:

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

# Ahora acepta diccionarios:
out = model(data.x_dict, data.edge_index_dict)

Usa (-1, -1) para canales de entrada bipartitos (fuente, destino pueden diferir). La inicialización perezosa maneja el resto.

2. Envoltorio HeteroConv — conv diferente por tipo de arista:

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

conv = HeteroConv({
    ('paper', 'cita', 'paper'): GCNConv(-1, 64),
    ('autor', 'escribe', 'paper'): SAGEConv((-1, -1), 64),
    ('paper', 'tiene', 'autor'): SAGEConv((-1, -1), 64),
})

El patrón de forward maneja tipos de arista ausentes y asegura que los parámetros compartidos se actualicen correctamente.

3. HGTConv (Transformer Heterogéneo) — una sola capa que maneja todos los tipos de arista automáticamente:

from torch_geometric.nn import HGTConv

class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, num_node_types, num_edge_types):
        super().__init__()
        self.lin_dict = torch.nn.ModuleDict()
        for node_type in num_node_types:
            self.lin_dict[node_type] = torch.nn.Linear(-1, hidden_channels)
        self.conv = HGTConv(hidden_channels, hidden_channels, num_heads, num_node_types, num_edge_types)
        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) 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()}

Cuándo usar cada uno:

• to_hetero(): Rápido para prototipado, cuando quieres probar la misma arquitectura homogénea en un grafo heterogéneo
• HeteroConv: Cuando necesitas diferentes capas conv o agregación por tipo de arista
• HGTConv: Para tareas complejas con muchos tipos de nodos/aristas donde la atención específica por tipo ayuda

Nota importante: los modelos heterogéneos requieren que especifiques el tipo de arista para la agregación. Pasa aggr='sum' o 'mean'/'max' en to_hetero() para definir cómo se combinan los mensajes de diferentes tipos de arista.

Lee references/heterogeneous.md para ejemplos completos de entrenamiento, integración con NeighborLoader y formas de exportar modelos heterogéneos a ONNX/TorchScript.

Guías de Referencia y Buenas Prácticas

Errores Comunes

1. Índices de esquina incorrectos: Siempre verifica que edge_index tenga forma [2, num_aristas] con dtype torch.long. Los pares de aristas como filas necesitan transposición.
2. Omitir funciones de activación: Las capas conv de PyG no aplican activación — añade ReLU, etc., entre capas.
3. Ignorar la máscara de entrenamiento: En clasificación de nodos, las divisiones train/val/test suelen estar predefinidas en el conjunto de datos (Cora, etc.). No mezcles datos de prueba en el entrenamiento.
4. Tamaño de lote en NeighborLoader: Usa batch[:batch.batch_size] para seleccionar solo los nodos semilla para la pérdida, no todos los nodos muestreados.
5. Falta de aristas inversas para grafos no dirigidos: Añade T.ToUndirected() en la transformación si tu modelo espera un grafo no dirigido.

Depuración de Formas de Tensor

Cada capa conv de PyG devuelve un tensor con la misma forma en la primera dimensión (número de nodos). La dimensión de características de salida es out_channels. Si las formas no coinciden, verifica:

• x.shape[0] == edge_index.max().item() + 1 (número de nodos consistente)
• Para capas sin peso compartido (RGCNConv, etc.), las dimensiones de entrada deben coincidir con los tipos de nodo

Manejo de Características de Arista

Muchas capas conv aceptan edge_attr. La forma esperada es [num_aristas, num_caracteristicas_arista]. Capas como GATConv y TransformerConv tienen soporte integrado; para otras, puede que necesites pasarlas manualmente en el método message().

Utilidades

# División de grafos en subgrafos
from torch_geometric.utils import to_undirected, add_self_loops, subgraph

# Convertir a no dirigido (ambas direcciones)
edge_index = to_undirected(edge_index)

# Añadir bucles propios
edge_index, _ = add_self_loops(edge_index, num_nodes=data.num_nodes)

# Extraer subgrafo
sub_edge_index, sub_nodes = subgraph(subset_indices, edge_index)

# Aleatorización y enlaces falsos
from torch_geometric.utils import erdos_renyi_graph, stochastic_blockmodel_graph

Salida de Modelo y Capas de Decisión

Para clasificación de nodos, la salida suelen ser logits por clase — aplica softmax en la inferencia si lo necesitas. Para predicción de enlaces, puntúa pares de nodos con un producto punto o una MLP. Para clasificación de grafos, después de varias capas conv, usa pooling global (mean/max/add) seguido de una capa lineal.

Lecturas Adicionales

Archivo	Contenido
references/message_passing.md	Implementaciones completas de GCN y EdgeConv, explicaciones detalladas de MessagePassing
references/link_prediction.md	Guía completa de predicción de enlaces: GAE/VGAE, bucles de entrenamiento, LinkNeighborLoader, evaluación
references/heterogeneous.md	Ejemplos completos de entrenamiento heterogéneo, integración con NeighborLoader, exportación
references/scaling.md	Entrenamiento multi-GPU, DDP, PyTorch Lightning, NeighborLoader, torch.compile, perfiles de memoria

Notas de Versión

PyG 2.7 introdujo cambios de ruptura:

• Python 3.9 y PyTorch ≤2.5 ya no son compatibles
• torch_geometric.distributed ha sido eliminado — usa torch.distributed
• torch_sparse ya no es necesario para operaciones dispersas básicas (PyG usa internamente torch.sparse)
• Se ha mejorado el soporte torch.compile para la mayoría de las capas conv
• Las nuevas capas incluyen GPSConv (para arquitecturas de transformador de grafo) y modelos mejorados de difusión

Consulta el changelog oficial para migraciones detalladas.

Documentación generada para torch-geometric 2.7.x, Oct 2025