Redes Neuronales de Grafos en PyTorch Sin Dolor Extra ni Workarounds Personalizados

La mayoría de los desarrolladores están acostumbrados a trabajar con estructuras de datos claras y regulares. Una imagen es una cuadrícula de píxeles y el texto encaja fácilmente en una secuencia de tokens. Las redes convolucionales regulares o los transformers manejan esto perfectamente.

Los problemas surgen cuando los datos son fundamentalmente no lineales. Tomemos las transacciones bancarias, las conexiones entre usuarios en redes sociales o las fórmulas químicas de moléculas complejas. Los grafos están presentes en todas partes. Si intentas alimentarlos a PyTorch regular, tendrás que luchar manualmente con matrices dispersas, escribir bucles complejos para la agregación de nodos y mantener el seguimiento de los índices.

Esto es exactamente para lo que se creó PyTorch Geometric (abreviado PyG). Es una extensión para PyTorch que maneja toda la matemática de bajo nivel por ti y proporciona una API intuitiva para trabajar con redes neuronales de grafos.

Cómo está Estructurado PyG

Si ya sabes cómo escribir modelos en PyTorch, te sentirás cómodo con PyG en un par de horas. El framework sigue los mismos principios: los mismos módulos torch.nn, bucles de entrenamiento familiares y trabajo directo con tensores.

Toda la magia de las redes de grafos se construye alrededor del concepto de Paso de Mensajes. Cada vértice en un grafo recoge información de sus vecinos, la combina y actualiza su propio estado.

PyG proporciona una clase base MessagePassing que abstrae esta lógica. Dentro, solo necesitas definir tres cosas:

• Cómo se forma un mensaje desde un nodo vecino
• Cómo se agregan estos mensajes (suma, media, máximo)
• Cómo se actualiza el propio nodo

Así es como se ve una implementación simple de una capa de convolución de grafos:

import torch
from torch.nn import Sequential, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr="max")
        self.mlp = Sequential(
            Linear(2 * in_channels, out_channels),
            ReLU(),
            Linear(out_channels, out_channels),
        )

    def forward(self, x, edge_index):
        # x задает фичи узлов, edge_index отвечает за связи между ними
        return self.propagate(edge_index, x=x)

    def message(self, x_j, x_i):
        # x_i — текущий узел, x_j — его сосед
        edge_features = torch.cat([x_i, x_j - x_i], dim=-1)
        return self.mlp(edge_features)

No necesitas escribir bucles sobre todos los bordes del grafo. PyG paraleliza esta operación por sí mismo y la ejecuta rápidamente gracias a los kernels CUDA compilados.

Ejemplo Básico: Clasificación de Nodos

Tomemos el dataset estándar Cora, que consiste en artículos científicos y las citas entre ellos. La tarea es predecir la categoría de un artículo basándose en su texto y sus conexiones con otras publicaciones.

Un modelo convolucional regular GCN con carga de datos se ensambla en literalmente un par de docenas de líneas:

import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv

# Загружаем датасет
dataset = Planetoid(root='.', name='Cora')

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 = self.conv2(x, edge_index)
        return x

model = GCN(dataset.num_features, 16, dataset.num_classes)

El bucle de entrenamiento se ve completamente estándar para los entusiastas de PyTorch: pasamos la matriz de características x y el tensor de bordes edge_index, calculamos CrossEntropy y llamamos a loss.backward().

El Principal Problema de los Grafos y Cómo PyG lo Resuelve

El principal desafío del aprendizaje profundo en grafos es el escalado. Las imágenes en un batch pueden dividirse fácilmente y enviarse a la GPU por partes. Pero en un grafo, todos los nodos están interconectados. Cuando un grafo crece hasta millones de nodos y miles de millones de bordes, ya no cabe en la memoria de incluso las GPUs más potentes.

Los desarrolladores de PyG dedicaron mucho esfuerzo a resolver este problema. La biblioteca contiene mecanismos de muestreo de subgrafos:

• NeighborLoader selecciona solo vecinos aleatorios para cada nodo en el batch, evitando que la memoria explote
• ClusterGCN agrupa un grafo grande en piezas independientes y entrena la red sobre ellas
• GraphSAINT muestrea subgrafos aleatorios mientras preserva su topología

Gracias a esto, PyG puede procesar grafos gigantes en GPUs regulares sin quedarse sin memoria.

Qué Más Hay Incluido de Fábrica

El repositorio contiene una impresionante colección de arquitecturas ya implementadas de artículos de investigación. Encontrarás GCN clásico, GAT y GraphSAGE, así como modelos especializados como SchNet y DimeNet para análisis molecular o PointNet para trabajar con nubes de puntos 3D.

Más allá de los algoritmos, la biblioteca incluye cargadores para cientos de datasets estándar: desde redes sociales hasta bioinformática. Esto ahorra mucho tiempo en escribir parsers.

Instalación y Piedras en el Camino

Desde la versión 2.3, el paquete base se instala extremadamente fácil:

pip install torch_geometric

Esto es suficiente para comenzar. Sin embargo, si necesitas máxima velocidad en grafos enormes, necesitarás instalar bibliotecas adicionales con extensiones C++/CUDA: pyg-lib, torch-scatter y torch-sparse.

Aquí es donde las cosas a veces se ponen difíciles. Estos binarios están fuertemente acoplados a versiones específicas de PyTorch y CUDA. Si instalas una compilación incompatible, Python te inundará con errores de importación de bibliotecas C++. Siempre verifica la matriz de compatibilidad en el sitio web del proyecto e instala las extensiones con URLs wheel explícitas.

Para Quién Es Este Framework

La biblioteca vale la pena probarla si tus datos no encajan bien en tablas o cuadrículas:

• Sistemas de recomendación: el grafo de relación "usuario-artículo" funciona mejor que la factorización de matrices regular.
• Anti-fraude y fintech: encontrar cadenas de transacciones sospechosas y grupos ocultos de estafadores.
• Bioinformática y química: predecir propiedades moleculares, descubrimiento de fármacos y análisis de proteínas.
• Procesamiento de datos 3D: los escaneos LiDAR y las nubes de puntos se representan perfectamente como grafos.

PyG se ha convertido en el estándar de facto en la industria para trabajar con grafos. Si planeas resolver tales problemas, definitivamente no vale la pena construir arquitecturas personalizadas desde cero — PyG te ahorrará mucho tiempo y recursos.