Classificazione dei grafi con i trasformatori

'Graph classification with transformers' in English.

Nel blog precedente, abbiamo esplorato alcuni aspetti teorici del machine learning sui grafi. Questo si concentrerà su come fare la classificazione dei grafi utilizzando la libreria Transformers. (Puoi anche seguirci scaricando il notebook di demo qui !)

Al momento, l’unico modello di trasformazione di grafi disponibile in Transformers è il Graphormer di Microsoft, quindi è quello che useremo qui. Non vediamo l’ora di vedere quali altri modelli le persone useranno e integreranno 🤗

Requisiti

Per seguire questo tutorial, è necessario avere installato datasets e transformers (versione >= 4.27.2), che è possibile fare con pip install -U datasets transformers.

Dati

Per utilizzare dati grafici, è possibile partire dai propri dataset o utilizzare quelli disponibili su Hub. Ci concentreremo sull’utilizzo di quelli già disponibili, ma sentiti libero di aggiungere i tuoi dataset!

Caricamento

Caricare un dataset grafico dall’Hub è molto semplice. Carichiamo il dataset ogbg-mohiv (una base di confronto dal Benchmark Open Graph di Stanford), archiviato nel repository OGB:

from datasets import load_dataset

# C'è solo una suddivisione sull'hub
dataset = load_dataset("OGB/ogbg-molhiv")

dataset = dataset.shuffle(seed=0)

Questo dataset ha già tre suddivisioni, train, validation e test, e tutte queste suddivisioni contengono le nostre 5 colonne di interesse (edge_index, edge_attr, y, num_nodes, node_feat), che puoi vedere eseguendo print(dataset).

Se hai altre librerie per i grafi, puoi usarle per tracciare i tuoi grafi e ispezionare ulteriormente il dataset. Ad esempio, utilizzando PyGeometric e matplotlib:

import networkx as nx
import matplotlib.pyplot as plt

# Vogliamo tracciare il primo grafo di train
graph = dataset["train"][0]

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

# Conversione in 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)])

# Tracciamento
nx.draw(G)

Formato

Sull’Hub, i dataset grafici sono principalmente archiviati come liste di grafi (utilizzando il formato jsonl).

Un singolo grafo è un dizionario, e questo è il formato atteso per i nostri dataset di classificazione dei grafi:

  • edge_index contiene gli indici dei nodi negli archi, archiviati come una lista contenente due liste parallele di indici di archi.
    • Tipo: lista di 2 liste di interi.
    • Esempio: un grafo contenente quattro nodi (0, 1, 2 e 3) e in cui le connessioni sono 1->2, 1->3 e 3->1 avrà edge_index = [[1, 1, 3], [2, 3, 1]]. Qui potresti notare che il nodo 0 non è presente qui, poiché non fa parte di un arco in senso stretto. Questo è il motivo per cui l’attributo successivo è importante.
  • num_nodes indica il numero totale di nodi disponibili nel grafo (per impostazione predefinita, si presume che i nodi siano numerati in sequenza).
    • Tipo: intero
    • Esempio: Nel nostro esempio precedente, num_nodes = 4.
  • y associa a ciascun grafo ciò che vogliamo prevedere da esso (sia una classe, un valore di proprietà o diverse etichette binarie per compiti diversi).
    • Tipo: lista di interi (per classificazione multi-classe), float (per regressione) o liste di uno e zero (per classificazione multi-task binaria)
    • Esempio: Potremmo prevedere la dimensione del grafo (piccola = 0, VoAGI = 1, grande = 2). Qui, y = [0].
  • node_feat contiene le caratteristiche disponibili (se presenti) per ogni nodo del grafo, ordinate per indice di nodo.
    • Tipo: lista di liste di interi (Opzionale)
    • Esempio: I nostri nodi sopra potrebbero avere, ad esempio, tipi (come diversi atomi in una molecola). Questo potrebbe dare node_feat = [[1], [0], [1], [1]].
  • edge_attr contiene gli attributi disponibili (se presenti) per ogni arco del grafo, seguendo l’ordinamento di edge_index.
    • Tipo: lista di liste di interi (Opzionale)
    • Esempio: I nostri archi sopra potrebbero avere, ad esempio, tipi (come legami molecolari). Questo potrebbe dare edge_attr = [[0], [1], [1]].

Preprocessing

I framework di trasformazione dei grafi di solito applicano una specifica operazione di preprocessing ai loro dataset per generare funzionalità e proprietà aggiuntive che aiutano l’attività di apprendimento sottostante (classificazione nel nostro caso). Qui, utilizziamo il preprocessing predefinito di Graphormer, che genera informazioni sul grado di ingresso/uscita, le matrici del percorso più breve tra i nodi e altre proprietà di interesse per il modello.

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

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

È anche possibile applicare questo preprocessing in tempo reale, nei parametri di DataCollator (impostando on_the_fly_processing su True): non tutti i dataset sono piccoli come ogbg-molhiv, e per i grafi grandi potrebbe essere troppo costoso memorizzare in anticipo tutti i dati preelaborati.

Modello

Caricamento

Qui, caricamo un modello/checkpoint preaddestrato esistente e lo affiniamo per il nostro compito successivo, che è una classificazione binaria (quindi num_classes = 2). Possiamo anche affinare il nostro modello per compiti di regressione (num_classes = 1) o per classificazione multi-task.

from transformers import GraphormerForGraphClassification

model = GraphormerForGraphClassification.from_pretrained(
    "clefourrier/pcqm4mv2_graphormer_base",
    num_classes=2, # num_classes per il compito successivo
    ignore_mismatched_sizes=True,
)

Esaminiamo questo in maggior dettaglio.

Chiamando il metodo from_pretrained sul nostro modello, i pesi vengono scaricati e memorizzati nella cache per noi. Poiché il numero di classi (per la predizione) dipende dal dataset, passiamo il nuovo valore di num_classes insieme a ignore_mismatched_sizes insieme a model_checkpoint. In questo modo viene creata una testa di classificazione personalizzata, specifica per il nostro compito, quindi probabilmente diversa dalla testa di decodifica originale.

È anche possibile creare un modello con inizializzazione casuale per addestrare da zero, seguendo i parametri noti di un determinato checkpoint o scegliendoli manualmente.

Allenamento o affinamento

Per addestrare il nostro modello in modo semplice, utilizzeremo un Trainer. Per istanziarlo, dovremo definire la configurazione di addestramento e la metrica di valutazione. La più importante è TrainingArguments, che è una classe che contiene tutti gli attributi per personalizzare l’addestramento. Richiede un nome di cartella, che verrà utilizzato per salvare i checkpoint del modello.

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, # la dimensione del batch può essere modificata automaticamente per evitare errori di memoria esaurita
    gradient_accumulation_steps=10,
    dataloader_num_workers=4, #1, 
    num_train_epochs=20,
    evaluation_strategy="epoch",
    logging_strategy="epoch",
    push_to_hub=False,
)

Per i dataset di grafi, è particolarmente importante sperimentare con le dimensioni dei batch e i passi di accumulo del gradiente per addestrare un numero sufficiente di campioni evitando errori di memoria esaurita.

L’ultimo argomento push_to_hub consente al Trainer di inviare regolarmente il modello all’Hub durante l’addestramento, ad ogni passaggio di salvataggio.

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

Nel Trainer per la classificazione dei grafi, è importante passare il collatore di dati specifico per il dataset di grafi dato, che converte i grafi individuali in batch per l’addestramento.

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

Quando il modello è addestrato, può essere salvato nell’hub con tutti gli artefatti di addestramento associati utilizzando push_to_hub.

Dato che questo modello è piuttosto grande, ci vogliono circa un giorno per addestrarlo/raffinarlo per 20 epoche su CPU (IntelCore i7). Per accelerare il processo, potresti utilizzare potenti GPU e la parallelizzazione, avviando il codice in un notebook Colab o direttamente nel cluster di tua scelta.

Nota finale

Ora che sai come utilizzare transformers per addestrare un modello di classificazione di grafi, speriamo che tu cercherai di condividere i tuoi checkpoint, modelli e dataset preferiti di trasformatori di grafi sul Hub per il resto della comunità possa utilizzarli!