Graphormer решает задачу классификации графов, адаптируя архитектуру трансформера под нерегулярные структуры данных. Модель от Microsoft доступна в библиотеке Transformers от Hugging Face и показывает сильные результаты на молекулярных датасетах. В этом руководстве разберём полный цикл: загрузку графового датасета, подготовку данных в формате edge_index, node_feat и edge_attr, настройку DataCollator и запуск обучения через Trainer.
Материал ориентирован на ML-инженеров и исследователей, которым нужен рабочий пример без длительного разбора документации. Вы получите готовый код, понимание внутренней механики модели и критерии оценки качества классификации.
Введение в Graphormer: трансформер для графов
Graphormer расширяет стандартный трансформер, добавляя структурные кодировки, которые передают модели информацию о топологии графа. Обычный трансформер принимает последовательность токенов, где позиция каждого элемента задана порядком. Граф не имеет естественного порядка узлов, поэтому позиционные эмбеддинги из классической архитектуры здесь не работают.
Модель решает проблему через три типа структурных кодировок: центральность узлов, пространственное отношение между парами узлов и кодировку рёбер. Эти сигналы добавляются к входным признакам и позволяют механизму внимания учитывать, какие узлы важнее и как они связаны.
Почему обычные трансформеры не подходят для графов?
Стандартный трансформер ожидает на вход матрицу размера [batch_size, seq_len, hidden_dim], где seq_len фиксирован. Графы в датасете имеют переменное число узлов и произвольные связи. Если просто вытянуть матрицу смежности в последовательность, модель потеряет информацию о соседстве и расстояниях между узлами.
Graphormer обходит это ограничение. Вместо позиционных эмбеддингов модель вычисляет центральность узла, например степень вершины, и использует её как сигнал важности. Для каждой пары узлов вычисляется кратчайший путь, который кодируется как пространственное отношение. Эти данные подаются в attention-механизм, что позволяет модели различать локальные и глобальные связи в графе.
Ключевые компоненты Graphormer: структурные кодировки
Центральность узла показывает, насколько вершина вовлечена в структуру графа. В реализации Graphormer используется степень узла, которая преобразуется в обучаемый эмбеддинг. Узел с высокой степенью получает отличное представление от узла с низкой степенью, что помогает модели учитывать топологическую значимость.
Пространственное отношение основано на кратчайшем пути между узлами. Если путь между узлами i и j имеет длину 3, это кодируется отдельным эмбеддингом. Такой подход позволяет attention-механизму учитывать, насколько далеко расположены узлы друг от друга, и по-разному взвешивать их взаимодействие.
Кодировка рёбер добавляет признаки связи между узлами. Для молекулярных графов это может быть тип химической связи: одинарная, двойная, ароматическая. Эти признаки суммируются с attention-скорами, что даёт модели информацию о характере взаимодействия.
Подготовка данных для Graphormer
Для обучения модели нужен датасет в формате, который понимает библиотека Transformers. Hugging Face Hub содержит готовые графовые датасеты, включая молекулярные бенчмарки. Разберём загрузку и структуру данных на примере ogbg-molhiv.
Загрузка датасета с Hugging Face Hub
Библиотека datasets позволяет загрузить графовый датасет одной командой:
from datasets import load_dataset
dataset = load_dataset("graphs-datasets/ogbg-molhiv")
print(dataset)Датасет разделён на train, validation и test. Каждый пример содержит словарь с полями edge_index, node_feat, edge_attr и labels. Просмотрим первый пример, чтобы понять структуру:
sample = dataset["train"][0]
print(sample.keys())
print(sample["edge_index"].shape)
print(sample["node_feat"].shape)
print(sample["edge_attr"].shape)Вывод покажет размерности тензоров. edge_index имеет форму [2, num_edges], node_feat - [num_nodes, num_node_features], edge_attr - [num_edges, num_edge_features]. Число узлов и рёбер различается для каждой молекулы, что требует специальной обработки при батчинге.
Формат данных: edge_index, node_feat, edge_attr
edge_index - это тензор размера 2 x num_edges. Первая строка содержит индексы исходных узлов, вторая - индексы целевых узлов. Для молекулы это связи между атомами. Если атом 0 соединён с атомом 1, в edge_index будет колонка [0, 1].
node_feat - матрица признаков узлов размера num_nodes x num_node_features. Для молекулярных графов признаки описывают свойства атомов: атомный номер, заряд, гибридизацию. Каждая строка соответствует одному атому.
edge_attr - матрица признаков рёбер размера num_edges x num_edge_features. Признаки описывают тип связи: одинарная, двойная, тройная, ароматическая. Эти данные помогают модели различать химические взаимодействия.
Для визуализации структуры можно вывести количество узлов и рёбер:
num_nodes = sample["node_feat"].shape[0]
num_edges = sample["edge_index"].shape[1]
print(f"Узлов: {num_nodes}, Рёбер: {num_edges}")Графы в датасете имеют разный размер, поэтому DataCollator должен уметь паддить их до одинаковой длины внутри батча.
Обучение Graphormer с помощью Trainer
После подготовки данных нужно загрузить модель, настроить DataCollator и запустить Trainer. Библиотека Transformers предоставляет готовый класс GraphormerForGraphClassification и GraphormerDataCollator.
Настройка DataCollator для графов
GraphormerDataCollator паддит графы до максимального размера в батче, добавляя маски для несуществующих узлов и рёбер. Это позволяет обрабатывать батчи с переменным числом узлов без потери информации:
from transformers import GraphormerDataCollator
collator = GraphormerDataCollator()
batch = collator([dataset["train"][i] for i in range(4)])
print(batch["attention_mask"].shape)Коллатор автоматически создаёт attention_mask, который указывает, какие позиции являются реальными узлами, а какие добавлены для паддинга. Модель игнорирует паддинговые позиции при вычислении внимания.
Запуск обучения и оценка результатов
Загрузим предобученную модель Graphormer для классификации графов. Для датасета ogbg-molhiv задача бинарная: предсказать, ингибирует ли молекула ВИЧ. Метрика качества - ROC-AUC:
from transformers import GraphormerForGraphClassification, Trainer, TrainingArguments
from sklearn.metrics import roc_auc_score
import numpy as np
model = GraphormerForGraphClassification.from_pretrained(
"clefourrier/graphormer-base-pcqm4mv2",
num_labels=2,
ignore_mismatched_sizes=True
)
def compute_metrics(eval_pred):
logits, labels = eval_pred
probs = 1 / (1 + np.exp(-logits[:, 1]))
auc = roc_auc_score(labels, probs)
return {"roc_auc": auc}
training_args = TrainingArguments(
output_dir="./graphormer-molhiv",
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
num_train_epochs=5,
evaluation_strategy="epoch",
save_strategy="epoch",
logging_steps=50,
learning_rate=2e-5,
weight_decay=0.01,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["validation"],
data_collator=collator,
compute_metrics=compute_metrics,
)
trainer.train()После обучения оцените модель на тестовом наборе. ROC-AUC выше 0.75 считается хорошим результатом для ogbg-molhiv без дополнительной настройки. Точное значение зависит от числа эпох и гиперпараметров.
Практический пример: классификация молекул
Полный цикл классификации молекул на датасете ogbg-molhiv занимает около 30 строк кода. Датасет содержит 41 127 молекул, каждая из которых размечена по свойству ингибирования ВИЧ. Задача бинарная, метрика - ROC-AUC.
Ключевые шаги:
- Загрузка датасета с Hugging Face Hub через load_dataset.
- Проверка формата данных: edge_index, node_feat, edge_attr.
- Создание GraphormerDataCollator для батчинга графов.
- Загрузка модели GraphormerForGraphClassification с num_labels=2.
- Настройка Trainer с compute_metrics для ROC-AUC.
- Запуск обучения на 5 эпохах с learning_rate 2e-5.
Результат на валидационном наборе обычно достигает ROC-AUC 0.78-0.82 после 5 эпох. Увеличение числа эпох до 10 может дать прирост 0.01-0.02 AUC, но повышает риск переобучения. Для улучшения результата можно использовать аугментацию графов или ансамблирование нескольких запусков.
Если вы работаете с другими графовыми задачами, например анализом социальных сетей или рекомендательными системами, формат данных остаётся тем же. Меняется только датасет и число классов в выходном слое.
Ограничения и соображения
Graphormer требует значительных вычислительных ресурсов. Обучение на датасете ogbg-molhiv с batch_size 16 занимает несколько часов на GPU с 16 ГБ памяти. Для больших графов с сотнями тысяч узлов потребуется распределённое обучение или уменьшение размера модели.
Модель чувствительна к качеству структурных кодировок. Если граф не имеет осмысленных признаков узлов или рёбер, центральность и кратчайшие пути не дадут достаточной информации для классификации. В таких случаях лучше использовать графовые нейронные сети, например GCN или GAT, которые напрямую агрегируют сообщения от соседей.
Для очень больших графов вычисление кратчайших путей между всеми парами узлов становится дорогим. Сложность preprocessing составляет O(N^2) по числу узлов, что ограничивает применение Graphormer на графах с более чем 10 000 узлов без предварительной кластеризации или сэмплирования.
Если ваша задача связана с анализом зависимостей в кодовой базе или оркестрацией AI-агентов, графовые подходы также применимы. Статья о GraphArc показывает, как графовая инженерия решает проблему непрозрачности агентных систем. Для управления ML-экспериментами и воспроизводимостью моделей полезно интегрировать обучение с MLflow.
Заключение
Graphormer - рабочий инструмент для классификации графов, доступный в экосистеме Hugging Face. Основные шаги: загрузка датасета, проверка формата edge_index, node_feat, edge_attr, настройка GraphormerDataCollator и запуск Trainer с метрикой ROC-AUC.
Модель показывает сильные результаты на молекулярных бенчмарках и может быть адаптирована под другие графовые задачи. Перед внедрением оцените размер графов и доступные вычислительные ресурсы. Для очень больших графов рассмотрите альтернативные подходы.
Документация Transformers содержит дополнительные примеры по графовым моделям. Начните с малого датасета, воспроизведите результат из этого руководства, затем масштабируйте на свои данные.