torch-geometric

PyTorch Geometric (PyG)

PyG est la bibliothèque standard pour les réseaux de neurones de graphes (GNN) construite sur PyTorch. Elle fournit des structures de données pour les graphes, plus de 60 implémentations de couches GNN, un entraînement mini-batch scalable et la prise en charge des graphes hétérogènes.

Installation

Testé avec torch-geometric 2.7.x (octobre 2025). Nécessite Python 3.10+ et PyTorch 2.6+.

# 1. Installez PyTorch d'abord (adaptez à votre configuration CUDA/CPU — voir https://pytorch.org/get-started/locally/)
uv pip install torch

# 2. PyG de base (aucun wheel d'extension requis pour une utilisation de base)
uv pip install torch_geometric

Les opérations accélérées optionnelles (pyg-lib, torch-scatter, torch-sparse, torch-cluster) ne sont pas nécessaires pour une utilisation de base de PyG (depuis PyG 2.3). Installez les wheels compatibles avec votre version depuis l'index de wheels PyG après avoir vérifié vos versions PyTorch et CUDA :

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# Ensuite, installez les wheels pour votre combo torch+CUDA, par exemple :
uv pip install pyg-lib torch-scatter torch-sparse torch-cluster \
  -f https://data.pyg.org/whl/torch-2.8.0+cu128.html

Vérifiez votre version :

import torch_geometric
print(torch_geometric.__version__)

Conda : le canal conda pyg n'est plus maintenu pour PyTorch >2.5 — utilisez uv pip install et l'index de wheels ci-dessus à la place.

Notes sur PyG 2.7

PyG 2.7 a abandonné la prise en charge de Python 3.9 et PyTorch ≤2.5. Consultez les notes de version 2.7.0 pour les tableaux de compatibilité PyTorch 2.6–2.8. torch_geometric.distributed est obsolète — utilisez le DDP standard torch.distributed (voir references/scaling.md).

Concepts Fondamentaux

Données de Graphe : Data et HeteroData

Un graphe est représenté par un objet Data. Les attributs clés :

from torch_geometric.data import Data

data = Data(
    x=node_features,          # [num_nodes, num_node_features]
    edge_index=edge_index,     # [2, num_edges] — format COO, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]
    y=labels,                  # niveau nœud [num_nodes, *] ou niveau graphe [1, *]
    pos=positions,             # [num_nodes, num_dimensions] (pour nuages de points/spatial)
)

Le format de edge_index est critique : c'est un tenseur [2, num_edges] où edge_index[0] = nœuds sources, edge_index[1] = nœuds cibles. Ce n'est PAS une liste de tuples. Si vous avez des paires d'arêtes sous forme de lignes, transposez et appelez .contiguous() :

# Si les arêtes sont [[src1, dst1], [src2, dst2], ...] — transposez d'abord :
edge_index = edge_pairs.t().contiguous()

Pour les graphes non orientés, incluez les deux directions : l'arête (0,1) nécessite à la fois [0,1] et [1,0] dans edge_index.

Pour les graphes hétérogènes, utilisez HeteroData — voir la section Graphes Hétérogènes ci-dessous.

Jeux de Données

PyG regroupe de nombreux jeux de données standard qui se téléchargent et se prétraitent automatiquement :

from torch_geometric.datasets import Planetoid, TUDataset

# Classification de nœuds sur un seul graphe (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]  # graphe unique avec masques train/val/test

# Classification multi-graphes (ENZYMES, MUTAG, IMDB-BINARY, etc.)
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... sont des graphes individuels

Jeux de données courants par tâche :

• Classification de nœuds : Planetoid (Cora/Citeseer/Pubmed), OGB (ogbn-arxiv, ogbn-products, ogbn-mag)
• Classification de graphes : TUDataset (MUTAG, ENZYMES, PROTEINS, IMDB-BINARY), OGB (ogbg-molhiv)
• Prédiction de liens : OGB (ogbl-collab, ogbl-citation2)
• Moléculaire : QM7, QM9, MoleculeNet
• Nuage de points/maillage : ShapeNet, ModelNet10/40, FAUST

Transformations

Les transformations prétraitent ou augmentent les données de graphe, analogues aux transformations torchvision :

import torch_geometric.transforms as T

# Transformations courantes
T.NormalizeFeatures()    # Normalisation des lignes des caractéristiques des nœuds pour que la somme soit 1
T.ToUndirected()         # Ajout des arêtes inverses pour rendre le graphe non orienté
T.AddSelfLoops()         # Ajout d'arêtes de boucle
T.KNNGraph(k=6)          # Construction d'un graphe k-NN à partir des positions d'un nuage de points
T.RandomJitter(0.01)     # Augmentation par bruit aléatoire sur les positions
T.Compose([...])         # Enchaînement de plusieurs transformations

# Appliquer en tant que pre_transform (une fois, sauvegardé sur disque) ou transform (à chaque accès)
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

Construction de Modèles GNN

Démarrage Rapide : Utilisation des Couches Intégrées

Le moyen le plus rapide de construire un GNN — empilez les couches de convolution depuis torch_geometric.nn :

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

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 = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

Important : les couches de convolution PyG n'incluent PAS de fonction d'activation — appliquez-les vous-même après chaque couche. C'est une décision de conception pour plus de flexibilité.

Choisir une Couche de Convolution

Choisissez en fonction de votre tâche et de la structure du graphe :

Couche	Meilleure pour	Idée clé
GCNConv	Classification de nœuds semi-supervisée homogène	Agrégation normalisée par degré, inspirée du spectral
GATConv / GATv2Conv	Lorsque l'importance des voisins varie	Messages pondérés par l'attention
SAGEConv	Grands graphes, contextes inductifs	Agrégation apprenable, compatible avec l'échantillonnage
GINConv	Classification de graphes, maximisation de l'expressivité	Aussi puissant que le test WL
TransformerConv	Caractéristiques d'arêtes riches, interactions complexes	Attention multi-têtes avec caractéristiques d'arêtes
EdgeConv	Nuages de points, graphes dynamiques	MLP sur les caractéristiques d'arêtes (x_i, x_j - x_i)
RGCNConv	Graphes hétérogènes avec de nombreux types de relations	Matrices de poids spécifiques aux relations
HGTConv	Graphes hétérogènes	Attention spécifique au type

Toutes les couches de convolution acceptent au minimum (x, edge_index). Beaucoup acceptent également edge_attr pour les caractéristiques d'arêtes.

Initialisation Paresseuse

Utilisez -1 pour les canaux d'entrée afin de laisser PyG déduire automatiquement les dimensions — particulièrement utile pour les modèles hétérogènes :

conv = SAGEConv((-1, -1), 64)  # Dimensions d'entrée déduites lors du premier passage forward
# Initialisez les modules paresseux :
with torch.no_grad():
    out = model(data.x, data.edge_index)

APIs de Modèle Haut Niveau

Pour les architectures courantes, PyG fournit des classes de modèle prêtes à l'emploi :

from torch_geometric.nn import GraphSAGE, GCN, GAT, GIN

model = GraphSAGE(
    in_channels=dataset.num_features,
    hidden_channels=64,
    out_channels=dataset.num_classes,
    num_layers=2,
)

Couches Personnalisées via MessagePassing

Pour implémenter une nouvelle couche GNN, sous-classez MessagePassing. Le cadre est :

1. propagate() orchestre le passage de messages
2. message() définit les informations qui circulent le long de chaque arête (la fonction phi)
3. aggregate() combine les messages à chaque nœud (somme/moyenne/max)
4. update() transforme le résultat agrégé (la fonction gamma)

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class MyConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # "add", "mean" ou "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # Prétraitement avant le passage de messages
        x = self.lin(x)
        # Début du passage de messages
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j : caractéristiques des nœuds sources pour chaque arête [num_edges, features]
        # Le suffixe _j indexe automatiquement les nœuds sources, _i indexe les nœuds cibles
        return x_j

La convention _i / _j : tout tenseur passé à propagate() peut être auto-indexé en ajoutant _i (nœud cible/central) ou _j (nœud source/voisin) dans la signature de message(). Ainsi, si vous passez x=... à propagate, vous pouvez accéder à x_i et x_j dans message().

Lisez references/message_passing.md pour les exemples d'implémentation complets de GCN et EdgeConv.

Schémas Spécifiques aux Tâches

Classification de Nœuds

# Entraînement complet sur un seul graphe (par exemple, Cora)
model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

# Évaluation — train(False) met le modèle en mode inférence (désactive le dropout/BN)
model.train(False)
pred = model(data.x, data.edge_index).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()

Classification de Graphes

Plusieurs graphes — utilisez DataLoader pour le mini-batch et le pooling global pour obtenir des représentations au niveau du graphe :

from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool

loader = DataLoader(dataset, batch_size=32, shuffle=True)

class GraphClassifier(torch.nn.Module):
    def __init__(self, in_ch, hidden_ch, out_ch):
        super().__init__()
        self.conv1 = GCNConv(in_ch, hidden_ch)
        self.conv2 = GCNConv(hidden_ch, hidden_ch)
        self.lin = torch.nn.Linear(hidden_ch, out_ch)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        x = global_mean_pool(x, batch)  # [num_graphs_in_batch, hidden_ch]
        return self.lin(x)

# Boucle d'entraînement
for data in loader:
    out = model(data.x, data.edge_index, data.batch)
    loss = F.cross_entropy(out, data.y)

Le DataLoader de PyG regroupe plusieurs graphes en créant des matrices d'adjacence bloc-diagonales. Le tenseur batch mappe chaque nœud à son index de graphe. Les opérations de pooling (global_mean_pool, global_max_pool, global_add_pool) l'utilisent pour agréger par graphe.

Prédiction de Liens

Divisez les arêtes en train/val/test, utilisez l'échantillonnage négatif :

from torch_geometric.transforms import RandomLinkSplit

transform = RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    is_undirected=True,
    add_negative_train_samples=False,
)
train_data, val_data, test_data = transform(data)

# Encodez les nœuds, puis scorez les arêtes
z = model.encode(train_data.x, train_data.edge_index)
# Arêtes positives
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)

Lisez references/link_prediction.md pour le guide complet de prédiction de liens : autoencodeurs GAE/VGAE, boucles d'entraînement complètes, LinkNeighborLoader pour les grands graphes, prédiction de liens hétérogènes et métriques d'évaluation.

Passage à l'Échelle pour les Grands Graphes

Pour les graphes qui ne tiennent pas dans la mémoire GPU, utilisez l'échantillonnage de voisins via NeighborLoader :

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[15, 10],     # Échantillonnez 15 voisins au saut 1, 10 au saut 2
    batch_size=128,              # Nombre de nœuds semences par lot
    input_nodes=data.train_mask, # Nœuds à partir desquels échantillonner
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # N'utilisez que les premiers nœuds batch_size pour la perte (ce sont les nœuds semences)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

Points clés sur NeighborLoader :

• La longueur de la liste num_neighbors doit correspondre à la profondeur du GNN (nombre de couches de passage de messages)
• Les nœuds semences sont toujours les premiers batch.batch_size nœuds dans la sortie
• batch.n_id mappe les indices réétiquetés aux ID de nœuds d'origine
• Fonctionne à la fois pour Data et HeteroData
• Pour la prédiction de liens, utilisez LinkNeighborLoader à la place
• L'échantillonnage de plus de 2-3 sauts est généralement irréalisable (explosion exponentielle)

Autres options de scalabilité : ClusterLoader (ClusterGCN), GraphSAINTSampler, ShaDowKHopSampler. Pour l'entraînement multi-GPU, DDP, l'intégration PyTorch Lightning et le support torch.compile, lisez references/scaling.md.

Graphes Hétérogènes

Pour les graphes avec plusieurs types de nœuds et d'arêtes (réseaux sociaux, graphes de connaissances, recommandation) :

from torch_geometric.data import HeteroData

data = HeteroData()

# Caractéristiques des nœuds — indexées par une chaîne de type de nœud
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)

# Indices d'arêtes — indexés par le triplet (type_src, type_arête, type_dst)
data['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))

# Dictionnaires de commodité d'accès
data.x_dict        # {'user': tensor, 'movie': tensor}
data.edge_index_dict  # {('user','rates','movie'): tensor, ...}
data.metadata()    # ([node_types], [edge_types])

Trois façons de construire des GNN hétérogènes

1. Conversion automatique avec to_hetero() — écrivez un modèle homogène, convertissez-le automatiquement :

from torch_geometric.nn import SAGEConv, to_hetero

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

model = GNN(64, dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

# Accepte maintenant des dictionnaires :
out = model(data.x_dict, data.edge_index_dict)

Utilisez (-1, -1) pour les canaux d'entrée bipartites (la source et la cible peuvent différer). L'initialisation paresseuse gère le reste.

2. Wrapper HeteroConv — une convolution différente par type d'arête :

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv

conv = HeteroConv({
    ('paper', 'cites', 'paper'): GCNConv(-1, 64),
    ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
    ('paper', 'published_in', 'venue'): GATConv((-1, -1), 64),
}, aggr='sum')

# Utilisation dans forward :
def forward(self, x_dict, edge_index_dict):
    x_dict = conv(x_dict, edge_index_dict)
    return x_dict

3. Couches hétérogènes dédiées — RGCNConv, HGTConv, HANConv pour les mécanismes d'attention/métaparcours :

from torch_geometric.nn import HGTConv

class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, metadata):
        super().__init__()
        self.lin_dict = torch.nn.ModuleDict()
        for node_type in metadata[0]:
            self.lin_dict[node_type] = torch.nn.Linear(-1, hidden_channels)
        self.conv = HGTConv(hidden_channels, hidden_channels, metadata, heads=4)
        self.out = torch.nn.Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        x_dict = {
            node_type: self.lin_dict[node_type](x).relu()
            for node_type, x in x_dict.items()
        }
        x_dict = self.conv(x_dict, edge_index_dict)
        return x_dict

Tous les modèles hétérogènes fonctionnent avec NeighborLoader acceptant l'argument input_nodes pour l'échantillonnage de voisins sur des données hétérogènes.

Optimisation des Performances

TorchScript et Compilation

PyG prend en charge torch.compile pour l'optimisation JIT :

model = torch.compile(model)  # Fusion GPU, quant à la compilation du noyau

Les couches PyG gérées (GCNConv, SAGEConv, etc.) sont toutes compatibles avec torch.compile. Pour les couches MessagePassing personnalisées, assurez-vous que les fonctions message(), aggregate() et update() sont traçables ou scriptables (pas de logique Python conditionnelle ou de dépendances externes dans le graphe de calcul).

Agrégation Personnalisée

Fournissez aggr_kwargs pour un réglage fin de l'agrégation :

conv = SAGEConv(64, 128, aggr='mean')
conv = SAGEConv(64, 128, aggr=MultiAggregation(['mean', 'max', 'std']))
conv = SAGEConv(64, 128, aggr=SoftmaxAggregation(learn=True))

Utilisation de la Mémoire

• Les mini-lots avec échantillonnage de voisins (NeighborLoader) réduisent considérablement la mémoire GPU
• Activez le mode mémoire torch.inference_mode() pendant l'évaluation si le rétropropagation n'est pas nécessaire
• Utilisez torch.set_float32_matmul_precision('high') (RTX 40xx/Ampere) ou torch.backends.cuda.matmul.allow_tf32 = True
• Pour les très grands graphes, GraphSaintSampler offre un échantillonnage informé du graphe entier avec des garanties de taille de mini-lot constantes

Lisez references/scaling.md pour les schémas avancés : apprentissage auto-supervisé, distillation de connaissances, quantification et plugiciels d'entraînement Lightning.

Dépannage

Problème	Solution
edge_index a la mauvaise forme	Devrait être [2, num_edges] ; transposer les paires d'arêtes sous forme de lignes edge_pairs.t().contiguous()
HeteroData ne fonctionne pas avec to_hetero	Assurez-vous que metadata() est correcte ; utilisez (-1, -1) pour les canaux d'entrée hétérogènes
DataLoader renvoie AttributeError pour batch	Vérifiez que DataLoader(batch_size=..., ...) est utilisé avec un jeu de données de graphes individuels, pas un seul Data
Erreur de version CUDA	Installez les wheels PyG depuis l'index des wheels correspondant à votre torch.version.cuda
L'échantillonnage de voisins est lent	Diminuez num_neighbors, limitez les sauts (hops) à ≤2, utilisez l'échantillonnage pyg-lib si disponible
Le module paresseux n'est pas initialisé	Passez des données factices via with torch.no_grad(): model(...) avant l'entraînement

Références Supplémentaires

Ces fichiers dans le même dossier contiennent des sujets avancés non abordés ici :

• references/message_passing.md — Guide MessagePassing complet : mise en œuvre complète de GCN et EdgeConv, fonctions d'agrégation personnalisées, fusion manuelle des arêtes, couches de pooling, modules d'attention et sérialisation.
• references/link_prediction.md — Guide complet de prédiction de liens : autoencodeurs GAE/VGAE, boucles d'entraînement, grand graphe LinkNeighborLoader, prédiction de liens hétérogènes, métriques de classement (MRR, Hits@K) et échantillonnage négatif.
• references/scaling.md — Scalabilité et performances : DDP multi-GPU, PyTorch Lightning, torch.compile, mélange de précision (FP16/BF16), fashion mini-batch, ClusterLoader, GraphSAINT et optimisation de la mémoire.

Consultez également :

• Site officiel de PyG
• Documentation PyG
• Dépôt GitHub PyG