Graph Neural Networks in PyTorch ohne zusätzliche Mühen und individuelle Workarounds

Die meisten Entwickler sind es gewohnt, mit klaren, regulären Datenstrukturen zu arbeiten. Ein Bild ist ein Raster aus Pixeln, und Text passt problemlos in eine Sequenz von Tokens. Reguläre Convolutional Networks oder Transformer bewältigen diese Aufgaben perfekt.

Probleme entstehen dort, wo Daten grundlegend nicht-linear sind. Man denke an Banktransaktionen, Verbindungen zwischen Nutzern in sozialen Netzwerken oder chemische Formeln komplexer Moleküle. Graphen sind überall im Einsatz. Wenn du versuchst, sie an reguläres PyTorch zu übergeben, musst du dich manuell mit Sparse-Matrizen herumschlagen, komplexe Schleifen für die Knotenaggregation schreiben und Indizes verfolgen.

Genau dafür wurde PyTorch Geometric (abgekürzt PyG) entwickelt. Es ist eine Erweiterung für PyTorch, die die gesamte Low-Level-Mathematik für dich erledigt und eine intuitive API für die Arbeit mit Graph Neural Networks bereitstellt.

Wie PyG aufgebaut ist

Wenn du bereits weißt, wie man Modelle in PyTorch schreibt, wirst du dich innerhalb weniger Stunden mit PyG vertraut machen. Das Framework folgt denselben Prinzipien: dieselben Module torch.nn, vertraute Trainingsschleifen und direkte Arbeit mit Tensoren.

Die gesamte Magie von Graph Networks basiert auf dem Message Passing-Konzept. Jeder Knoten in einem Graphen sammelt Informationen von seinen Nachbarn, kombiniert sie und aktualisiert seinen eigenen Zustand.

PyG bietet eine Basisklasse MessagePassing, die diese Logik abstrahiert. Darin musst du nur drei Dinge definieren:

• Wie eine Nachricht von einem benachbarten Knoten gebildet wird
• Wie diese Nachrichten aggregiert werden (Summe, Mittelwert, Maximum)
• Wie der Knoten selbst aktualisiert wird

So sieht eine einfache Graph-Convolution-Layer-Implementierung aus:

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)

Du musst keine Schleifen über alle Graph-Kanten schreiben. PyG parallelisiert diese Operation selbst und führt sie dank kompilierter CUDA-Kernels schnell aus.

Grundlegendes Beispiel: Knotenklassifikation

Nehmen wir den Standard-Datensatz Cora, der aus wissenschaftlichen Arbeiten und Zitationen zwischen ihnen besteht. Die Aufgabe besteht darin, die Kategorie einer Arbeit basierend auf ihrem Text und ihren Verbindungen zu anderen Publikationen vorherzusagen.

Ein reguläres Convolutional-Modell GCN mit Datenladen wird buchstäblich in knapp zwei Dutzend Zeilen zusammengestellt:

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)

Die Trainingsschleife sieht für PyTorch-Enthusiasten völlig standardmäßig aus: wir übergeben die Feature-Matrix x und den Edge-Tensor edge_index, berechnen CrossEntropy und rufen loss.backward() auf.

Das Hauptproblem von Graphen und wie PyG es löst

Die Hauptherausforderung beim Deep Learning auf Graphen ist die Skalierung. Bilder in einem Batch können leicht aufgeteilt und in Teilen an die GPU gesendet werden. Aber in einem Graphen sind alle Knoten miteinander verbunden. Wenn ein Graph auf Millionen von Knoten und Milliarden von Kanten anwächst, passt er nicht mehr in den Speicher selbst der leistungsstärksten GPUs.

Die PyG-Entwickler haben viel Aufwand in die Lösung dieses Problems gesteckt. Die Bibliothek enthält Subgraph-Sampling-Mechanismen:

• NeighborLoader wählt nur zufällige Nachbarn für jeden Knoten im Batch aus und verhindert, dass der Speicher explodiert
• ClusterGCN clustert einen großen Graphen in unabhängige Stücke und trainiert das Netzwerk darauf
• GraphSAINT sampelt zufällige Subgraphen unter Beibehaltung ihrer Topologie

Dadurch kann PyG riesige Graphen auf normalen GPUs verarbeiten, ohne out-of-memory zu gehen.

Was sonst noch out of the box enthalten ist

Das Repository enthält eine beeindruckende Sammlung bereits implementierter Architekturen aus Forschungsarbeiten. Du findest klassisches GCN, GAT und GraphSAGE sowie spezialisierte Modelle wie SchNet und DimeNet für molekulare Analysen oder PointNet für die Arbeit mit 3D-Punktwolken.

Über Algorithmen hinaus enthält die Bibliothek Loader für Hunderte von Standard-Datensätzen: von sozialen Netzwerken bis zur Bioinformatik. Das spart viel Zeit beim Schreiben von Parsern.

Installation und Fallstricke

Seit Version 2.3 lässt sich das Basispaket extrem einfach installieren:

pip install torch_geometric

Das reicht aus, um loszulegen. Wenn du jedoch maximale Geschwindigkeit bei riesigen Graphen benötigst, musst du zusätzliche Bibliotheken mit C++/CUDA-Erweiterungen installieren: pyg-lib, torch-scatter und torch-sparse.

Hier wird es manchmal knifflig. Diese Binärdateien sind eng an bestimmte Versionen von PyTorch und CUDA gekoppelt. Wenn du eine inkompatible Version installierst, wird Python dich mit C++-Bibliotheks-Importfehlern überschütten. Überprüfe immer die Kompatibilitätsmatrix auf der Projektwebsite und installiere Erweiterungen mit expliziten Wheel-URLs.

Für wen dieses Framework geeignet ist

Die Bibliothek lohnt sich, wenn deine Daten nicht gut in Tabellen oder Raster passen:

• Empfehlungssysteme: der „User-Item-Beziehungsgraph" funktioniert besser als reguläre Matrixfaktorisierung.
• Betrugsbekämpfung und Fintech: Aufdeckung von Ketten verdächtiger Transaktionen und versteckter Gruppen von Betrügern.
• Bioinformatik und Chemie: Vorhersage molekularer Eigenschaften, Drug Discovery und Proteinanalysen.
• 3D-Datenverarbeitung: LiDAR-Scans und Punktwolken werden perfekt als Graphen dargestellt.

PyG hat sich zum De-facto-Standard in der Branche für die Arbeit mit Graphen entwickelt. Wenn du planst, solche Probleme zu lösen, lohnt es sich definitiv nicht, individuelle Architekturen von Grund auf neu zu bauen – PyG spart dir viel Zeit und Ressourcen.