alterlab-torch-geometric

PyTorch Geometric (PyG)

Overview

PyTorch Geometric is a library built on PyTorch for developing and training Graph Neural Networks (GNNs). Apply this skill for deep learning on graphs and irregular structures, including mini-batch processing, multi-GPU training, and geometric deep learning applications.

When to Use This Skill

This skill should be used when working with:

• Graph-based machine learning: Node classification, graph classification, link prediction
• Molecular property prediction: Drug discovery, chemical property prediction
• Social network analysis: Community detection, influence prediction
• Citation networks: Paper classification, recommendation systems
• 3D geometric data: Point clouds, meshes, molecular structures
• Heterogeneous graphs: Multi-type nodes and edges (e.g., knowledge graphs)
• Large-scale graph learning: Neighbor sampling, distributed training

Quick Start

uv pip install torch_geometric

Graphs are torch_geometric.data.Data objects: x (node features [N, F]), edge_index (connectivity in COO [2, E]), optional edge_attr, y, pos, and any custom attribute (train_mask, etc.). DataLoader batches multiple graphs into one block-diagonal graph (no padding); a batch vector maps nodes back to their source graph.

Full install/sparse-deps, basic graph creation, benchmark loading, edge-index format, and mini-batching details: references/getting_started.md.

Core Workflow

1. Load or build data — benchmark datasets, custom InMemoryDataset, or from CSV (references/datasets_and_loading.md; full catalog in references/datasets_reference.md).
2. Define a GNN — stack pre-built conv layers (GCNConv, GATConv, SAGEConv) or subclass MessagePassing for custom layers (references/building_gnns.md; full layer list in references/layers_reference.md).
3. Train — node classification (single graph, train/test masks), graph classification (DataLoader + global pooling), or large-scale via NeighborLoader neighbor sampling (references/training_workflows.md).
4. Go advanced if needed — HeteroData/to_hetero for heterogeneous graphs, transforms, GNNExplainer explainability, hierarchical pooling, GPU, save/load (references/advanced_features.md; transforms catalog in references/transforms_reference.md).