Facebook представила фреймворк PyTorch BigGraph

На конференции SysML компания Facebook представила PyTorch BigGraph (PBG) — распределенную систему для обучения векторного представления графов. Особенность этой системы в том, что она рассчитана на большие графы, содержащие до миллиардов вершин и триллионов ребер.

Документация по фреймворку находится в открытом доступе.

Векторное представление графов — одна из форм обучения без учителя, которая используется для задач прогнозирования в социальных сетях, обнаружения паттернов в IoT, моделирования последовательностей приема лекарств и др. В таких случаях данные естественным образом представляются с использованием графовых структур.

PBG позволяет быстрее создавать векторные представления графовых структур для больших графов в моделях PyTorch. Метод обучения векторных представлений для графов похож на аналогичный для текста — word2vec. Обучение происходит таким образом, чтобы пары вершин, между которыми есть ребра, находились друг к другу ближе, чем не соединенные ребром вершины.

Проблема

Большинство методов построения векторных представлений не подходят для применения к очень большим графовым структурам. К примеру, для модели с двумя миллиардами узлов и сотней параметров векторного представления на узел (выраженных числами с плавающей точкой), потребовалось 800 Гб памяти только для хранения параметров, поэтому требования стандартных методов превышают объем памяти типичных облачных серверов. Для решения этой проблемы и был разработан BigGraph.

Принцип работы

PBG использует четыре подхода:

• Разбиение графа. Теперь модель не потребуется выгружать в память полностью.
• Многопоточные вычисления на каждой машине.
• Распределенное выполнение на нескольких машинах (необязательно), все из них работают над разными частями графа.
• Порционное создание выборки отрицательных ребер для каждого положительного ребра. Позволяет обрабатывать > 1 миллиона ребер/сек на каждой машине.

PBG произвольно делит структуру графа на P разделов с таким объемом, чтобы два раздела могли поместиться в памяти.embedding

Если ребро (edge) имеет начальную вершину (node) в разделе p1 и конечную вершину в разделе p2, то он помещается в ячейку (bucket) (p1, p2). Затем ребра графа разбиваются на P2 ячеек в зависимости от их начальной и конечной вершин. После разделения вершин и ребер обучение может выполняться для одной ячейки за раз. Обучение ячейки (p1, p2) требует только того, чтобы векторные представления для разделов p1 и p2 были сохранены в памяти. Структура PBG гарантирует, что у ячеек есть обученное ранее векторное представление хотя бы для одного из разделов.

PBG является одним из первых методов, которые могут масштабироваться и обрабатывать данные графовых структур с миллиардами узлов и триллионами ребер, что позволяет ожидать интересных результатов в ближайшем будущем.