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).

Building GNNs (at a glance)

GNNs follow neighborhood aggregation: transform node features → propagate messages along edges → aggregate from neighbors → update representations. PyG ships 40+ conv layers. When choosing one, check its capabilities: SparseTensor support, edge_weight, edge_attr, bipartite, and lazy (-1 channel) initialization. Code for GCN/GAT/GraphSAGE and custom MessagePassing layers (including the _i/_j target/source naming convention) is in references/building_gnns.md.

Resources

Bundled References

This skill includes detailed reference documentation:

• references/getting_started.md: Install, basic graph creation, Data structure, edge-index format, mini-batching
• references/building_gnns.md: Message passing, GCN/GAT/GraphSAGE code, custom MessagePassing layers, layer capabilities
• references/datasets_and_loading.md: Built-in datasets, custom InMemoryDataset, loading graphs from CSV
• references/training_workflows.md: Node classification, graph classification, large-scale neighbor sampling
• references/advanced_features.md: Heterogeneous graphs, transforms, explainability, pooling, GPU, save/load
• references/layers_reference.md: Complete listing of all 40+ GNN layers with descriptions and capabilities
• references/datasets_reference.md: Comprehensive dataset catalog organized by category
• references/transforms_reference.md: All available transforms and their use cases

Scripts

Utility scripts are provided in scripts/:

• scripts/visualize_graph.py: Visualize graph structure using networkx and matplotlib
• scripts/create_gnn_template.py: Generate boilerplate code for common GNN architectures
• scripts/benchmark_model.py: Benchmark model performance on standard datasets

Execute scripts directly or read them for implementation patterns.

Official Resources

• Documentation: https://pytorch-geometric.readthedocs.io/
• GitHub: https://github.com/pyg-team/pytorch_geometric
• Tutorials: https://pytorch-geometric.readthedocs.io/en/latest/get_started/introduction.html
• Examples: https://github.com/pyg-team/pytorch_geometric/tree/master/examples