Clasificación de gráficos con Transformers

Graph Classification with Transformers

En el blog anterior, exploramos algunos aspectos teóricos del aprendizaje automático en grafos. Este se centrará en cómo realizar la clasificación de grafos utilizando la biblioteca Transformers. (¡También puedes seguirlo descargando el cuaderno de demostración aquí!)

Actualmente, el único modelo de transformador de gráficos disponible en Transformers es Graphormer de Microsoft, por lo que es el que utilizaremos aquí. Estamos ansiosos por ver qué otros modelos utilizarán e integrarán 🤗

Requisitos

Para seguir este tutorial, necesitas tener instalados los paquetes datasets y transformers (versión >= 4.27.2), lo cual puedes hacer con el comando pip install -U datasets transformers.

Datos

Para utilizar datos de grafos, puedes comenzar desde tus propios conjuntos de datos o utilizar los disponibles en el Hub. Nos centraremos en utilizar los que ya están disponibles, ¡pero siéntete libre de agregar tus propios conjuntos de datos!

Carga

Cargar un conjunto de datos de gráficos desde el Hub es muy fácil. Carguemos el conjunto de datos ogbg-mohiv (un referencial del Open Graph Benchmark de Stanford), almacenado en el repositorio OGB:

• Cómo alojar un juego de Unity en un Espacio
• Presentando el blog de HuggingFace para hablantes de chino Fomentan...
• Entrenando un modelo de lenguaje con 🤗 Transformers utilizando Ten...

from datasets import load_dataset

# Solo hay una división en el Hub
dataset = load_dataset("OGB/ogbg-molhiv")

dataset = dataset.shuffle(seed=0)

Este conjunto de datos ya tiene tres divisiones: train, validation y test, y todas estas divisiones contienen nuestras 5 columnas de interés (edge_index, edge_attr, y, num_nodes, node_feat), las cuales puedes ver ejecutando print(dataset).

Si tienes otras bibliotecas de gráficos, puedes utilizarlas para representar gráficos y examinar más detenidamente el conjunto de datos. Por ejemplo, utilizando PyGeometric y matplotlib:

import networkx as nx
import matplotlib.pyplot as plt

# Queremos representar el primer gráfico de entrenamiento
graph = dataset["train"][0]

edges = graph["edge_index"]
num_edges = len(edges[0])
num_nodes = graph["num_nodes"]

# Conversión al formato networkx
G = nx.Graph()
G.add_nodes_from(range(num_nodes))
G.add_edges_from([(edges[0][i], edges[1][i]) for i in range(num_edges)])

# Representación gráfica
nx.draw(G)

Formato

En el Hub, los conjuntos de datos de gráficos se almacenan principalmente como listas de gráficos (utilizando el formato jsonl).

Un solo gráfico es un diccionario, y aquí está el formato esperado para nuestros conjuntos de datos de clasificación de gráficos:

• edge_index contiene los índices de los nodos en los bordes, almacenados como una lista que contiene dos listas paralelas de índices de bordes.
  • Tipo: lista de 2 listas de enteros.
  • Ejemplo: un gráfico que contiene cuatro nodos (0, 1, 2 y 3) y donde las conexiones son 1->2, 1->3 y 3->1 tendrá edge_index = [[1, 1, 3], [2, 3, 1]]. Aquí puedes notar que el nodo 0 no está presente, ya que no forma parte de un borde en sí. Por eso es importante el siguiente atributo.
• num_nodes indica el número total de nodos disponibles en el gráfico (por defecto, se asume que los nodos están numerados secuencialmente).
  • Tipo: entero
  • Ejemplo: En nuestro ejemplo anterior, num_nodes = 4.
• y asigna a cada gráfico lo que queremos predecir a partir de él (ya sea una clase, un valor de propiedad o varias etiquetas binarias para diferentes tareas).
  • Tipo: lista de enteros (para clasificación multiclase), flotantes (para regresión) o listas de unos y ceros (para clasificación binaria multitarea)
  • Ejemplo: Podríamos predecir el tamaño del gráfico (pequeño = 0, VoAGI = 1, grande = 2). Aquí, y = [0].
• node_feat contiene las características disponibles (si están presentes) para cada nodo del gráfico, ordenadas por índice de nodo.
  • Tipo: lista de listas de enteros (opcional)
  • Ejemplo: Nuestros nodos anteriores podrían tener, por ejemplo, tipos (como diferentes átomos en una molécula). Esto podría dar node_feat = [[1], [0], [1], [1]].
• edge_attr contiene los atributos disponibles (si están presentes) para cada borde del gráfico, siguiendo el orden de edge_index.
  • Tipo: lista de listas de enteros (opcional)
  • Ejemplo: Nuestros bordes anteriores podrían tener, por ejemplo, tipos (como enlaces moleculares). Esto podría dar edge_attr = [[0], [1], [1]].

Preprocesamiento

Los marcos de trabajo para transformar gráficos generalmente aplican un preprocesamiento específico a sus conjuntos de datos para generar características y propiedades adicionales que ayudan a la tarea de aprendizaje subyacente (clasificación en nuestro caso). Aquí, utilizamos el preprocesamiento predeterminado de Graphormer, que genera información de grado de entrada/salida, la ruta más corta entre matrices de nodos y otras propiedades de interés para el modelo.

from transformers.models.graphormer.collating_graphormer import preprocess_item, GraphormerDataCollator

dataset_processed = dataset.map(preprocess_item, batched=False)

También es posible aplicar este preprocesamiento sobre la marcha, en los parámetros del DataCollator (configurando on_the_fly_processing en True): no todos los conjuntos de datos son tan pequeños como ogbg-molhiv, y para gráficos grandes, podría ser demasiado costoso almacenar todos los datos preprocesados de antemano.

Modelo

Carga

Aquí, cargamos un modelo/checkpoint pre-entrenado existente y lo afinamos en nuestra tarea secundaria, que es una tarea de clasificación binaria (por lo tanto, num_classes = 2). También podríamos afinar nuestro modelo en tareas de regresión (num_classes = 1) o en clasificación de múltiples tareas.

from transformers import GraphormerForGraphClassification

model = GraphormerForGraphClassification.from_pretrained(
    "clefourrier/pcqm4mv2_graphormer_base",
    num_classes=2, # num_classes para la tarea secundaria
    ignore_mismatched_sizes=True,
)

Veámoslo con más detalle.

Llamar al método from_pretrained en nuestro modelo descarga y almacena en caché los pesos para nosotros. Como el número de clases (para la predicción) depende del conjunto de datos, pasamos el nuevo num_classes así como ignore_mismatched_sizes junto con model_checkpoint. Esto asegura que se cree una cabeza de clasificación personalizada, específica para nuestra tarea, por lo tanto, probablemente diferente de la cabeza decodificadora original.

También es posible crear un modelo nuevo inicializado aleatoriamente para entrenar desde cero, siguiendo los parámetros conocidos de un checkpoint dado o eligiéndolos manualmente.

Entrenamiento o afinamiento

Para entrenar nuestro modelo de manera sencilla, utilizaremos un Trainer. Para instanciarlo, necesitaremos definir la configuración de entrenamiento y la métrica de evaluación. Lo más importante es TrainingArguments, que es una clase que contiene todos los atributos para personalizar el entrenamiento. Requiere un nombre de carpeta, que se utilizará para guardar los puntos de control del modelo.

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    "graph-classification",
    logging_dir="graph-classification",
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    auto_find_batch_size=True, # el tamaño del lote puede cambiarse automáticamente para evitar errores de falta de memoria
    gradient_accumulation_steps=10,
    dataloader_num_workers=4, #1, 
    num_train_epochs=20,
    evaluation_strategy="epoch",
    logging_strategy="epoch",
    push_to_hub=False,
)

Para conjuntos de datos gráficos, es particularmente importante probar diferentes tamaños de lotes y pasos de acumulación de gradientes para entrenar en suficientes muestras y evitar errores de falta de memoria.

El último argumento push_to_hub permite que el Trainer envíe el modelo al Hub regularmente durante el entrenamiento, en cada paso de guardado.

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset_processed["train"],
    eval_dataset=dataset_processed["validation"],
    data_collator=GraphormerDataCollator(),
)

En el Trainer para clasificación de gráficos, es importante pasar el recolector de datos específico para el conjunto de datos gráfico dado, que convertirá gráficos individuales en lotes para el entrenamiento.

train_results = trainer.train()
trainer.push_to_hub()

Cuando el modelo está entrenado, se puede guardar en el hub con todos los artefactos de entrenamiento asociados utilizando push_to_hub.

Dado que este modelo es bastante grande, tarda aproximadamente un día en entrenarse/afinarse durante 20 épocas en CPU (IntelCore i7). Para acelerar el proceso, podrías usar GPU potentes y paralelización en su lugar, lanzando el código en un cuaderno de Colab o directamente en el clúster de tu elección.

Nota final

Ahora que sabes cómo usar transformers para entrenar un modelo de clasificación de gráficos, ¡esperamos que intentes compartir tus puntos de control, modelos y conjuntos de datos de transformers favoritos en el Hub para que el resto de la comunidad los use!