DOCSEARCHbyAlgolia
Submit Docs

PyTorch Geometric (PyG)

repository·master·Indexed 12 days ago

https://github.com/pyg-team/pytorch_geometric

A specialized library for deep learning on graphs and other irregular structures (geometric deep learning) that extends PyTorch. Version 2.9.0 provides efficient message-passing layers, large-scale graph sampling, and support for heterogeneous and dynamic graphs.

Tokens
71.9K
Snippets
163
Records
275
Agent score
96%

What's inside PyG

  1. Overview of PyTorch Geometric (PyG)

    master

    PyTorch Geometric (PyG) is a library built on top of PyTorch designed for writing and training Graph Neural Networks (GNNs) on structured data and irregular structures (geometric deep learning).

    Key capabilities include:

    • Deep Learning on Graphs: Implementations of methods from various published papers for graphs, 3D meshes, and point clouds.
    • Data Handling: Mini-batch loaders for both many small graphs and single giant graphs, and support for DataPipe.
    • Performance & Scaling: Support for multi-GPU training, torch.compile, and distributed training.
    • Datasets & Transforms: A large collection of common benchmark datasets and helpful transforms for graph and geometric data.
    • Advanced Features: Support for sparse tensors, JIT, remote execution, and GraphGym.

  2. Overview of Distributed Training in PyG

    master

    PyG provides an in-house distributed training solution (available from version 2.5 onwards) designed for large-scale datasets that cannot fit into the memory of a single machine.

    This architecture enables training on graphs with billions of nodes by partitioning the graph across a cluster of CPUs. It combines PyTorch's Distributed Data Parallel (DDP) for model training with Remote Procedure Calls (RPCs) for efficient sampling and retrieval of non-local features. This allows for synchronized model training across multiple nodes using the TCP/IP protocol and the gloo communication backend.

  3. Overview of LLM and GNN Co-training Examples

    master

    PyTorch Geometric provides several example implementations for co-training Large Language Models (LLMs) and Graph Neural Networks (GNNs). Key examples include:

    • g_retriever.py: Helper functions for the G-retriever module. Includes Neo4j integration support.
    • txt2kg_rag.py: A complete end-to-end RAG (Retrieval-Augmented Generation) pipeline using TXT2KG, Vector, and Graph RAG with a GNN.
    • txt2qa.py: A synthetic multi-hop QA generation pipeline from text documents. Supports both vLLM (local GPU) and NVIDIA NIM (API) backends.
    • molecule_gpt.py: Implementation for MoleculeGPT for molecular property prediction.
    • glem.py: Implementation for GLEM using a variational EM framework.
    • git_mol.py: Implementation for GIT-Mol (multi-modal LLM for molecular science).
    • protein_mpnn.py: Implementation for ProteinMPNN protein sequence design.

  4. Use GNN operators and utilities for augmentation and normalization

    master

    PyG includes a rich set of neural network operators and graph utilities designed to enhance model performance through augmentation, regularization, and normalization. These can be applied to existing GNN layers.

    Augmentation and Regularization

    • DropEdge: Randomly removes edges to prevent overfitting.
    • DropNode: Randomly removes nodes.
    • MaskFeature: Masks specific node features.
    • AddRandomEdge: Adds random edges to the graph.
    • DropPath: Used in masked graph modeling.
    • ShuffleNode: Shuffles node order for contrastive learning.

    Normalization and Transforms

    • GraphNorm: A principled approach to accelerate GNN training.
    • GDC (Graph Diffusion Convolution): Uses diffusion to improve graph learning via transforms.

  5. Explore torch_geometric.datasets for graph data

    master

    The torch_geometric.datasets module provides a wide variety of dataset loaders and generators for different graph types. You can use these to quickly load standard benchmark datasets or generate synthetic graphs for testing and experimentation.

    Datasets are categorized by their structural properties:

    • Homogeneous Datasets: Standard graphs where all nodes and edges share the same type.
    • Heterogeneous Datasets: Graphs containing multiple types of nodes and edges (e.g., user-product interaction graphs).
    • Hypergraph Datasets: Datasets where edges can connect more than two nodes.
    • Synthetic Datasets: Pre-defined synthetic graph datasets used for benchmarking.
    • Graph Generators: Tools to generate custom graph structures (found in torch_geometric.datasets.graph_generator).
    • Motif Generators: Tools specifically for generating graphs based on specific motifs (found in torch_geometric.datasets.motif_generator).

  7. Explore experimental features in torch_geometric.contrib

    master

    The torch_geometric.contrib package contains several sub-modules for experimental graph machine learning components:

    • Convolutional Layers (torch_geometric.contrib.nn.conv): Experimental graph convolutional layer implementations.
    • Models (torch_geometric.contrib.nn.models): Experimental end-to-end GNN model architectures.
    • Datasets (torch_geometric.contrib.datasets): Experimental graph datasets.
    • Transforms (torch_geometric.contrib.transforms): Experimental data transformation utilities for graph data.
    • Explainer (torch_geometric.contrib.explain): Experimental tools for explaining GNN predictions.

  8. Explore torch_geometric.data modules

    master

    The torch_geometric.data module provides the core abstractions for handling graph data in PyTorch Geometric. It is organized into several functional areas:

    • Data Objects: Core classes for representing graphs and geometric data (e.g., Data, HeteroData).
    • Remote Backend Interfaces: Classes for interacting with remote data backends.
    • Databases: Interfaces for managing graph data stored in databases.
    • PyTorch Lightning Wrappers: Specialized wrappers located in torch_geometric.data.lightning to facilitate integration with the PyTorch Lightning framework.
    • Helper Functions: Utility functions for data manipulation and processing.

  9. Use the PyG Benchmark Suite for evaluation

    master

    The PyG Benchmark Suite provides standardized evaluation scripts to compare different methods in homogeneous scenarios. It is designed to avoid hyperparameter and model selection on the test set by utilizing an additional validation set.

    The suite includes evaluation scripts for the following tasks:

    • Semi-supervised node classification: Located in benchmark/citation.
    • Graph classification: Located in benchmark/kernel.
    • Point cloud classification: Located in benchmark/points.
    • Runtimes: Located in benchmark/runtime.

  10. Explore available Graph Convolutional layers

    master

    PyTorch Geometric provides a wide variety of specialized graph convolution layers based on recent research. Key implementations include:

    • GMMConv: Geometric Deep Learning on Graphs and Manifolds using Mixture Model CNNs.
    • FeaStConv: Feature-Steered Graph Convolutions for 3D Shape Analysis.
    • PointTransformerConv: Point Transformer architecture.
    • HypergraphConv: Hypergraph Convolution and Hypergraph Attention.
    • GravNetConv: Distance-weighted Graph Networks for irregular particle-detector geometry.
    • SuperGAT: Graph Attention Design With Self-Supervision.
    • HGTConv: Heterogeneous Graph Transformer.
    • HEATConv: Heterogeneous Edge-Enhanced Graph Attention Network.
    • SSGConv: Simple Spectral Graph Convolution.
    • FusedGATConv: Coordinated Computation, IO, and Memory perspective for GNNs.
    • GPSConv: A general, powerful, and scalable Graph Transformer recipe.

  11. Explore implemented GNN models and architectures

    master

    PyTorch Geometric provides a wide range of pre-implemented Graph Neural Network (GNN) models, ranging from specialized architectures to general-purpose frameworks.

    Specialized Architectures

    • JumpingKnowledge: For representation learning with varying neighborhood scales.
    • MetaLayer: A flexible framework for building custom graph networks (similar to DeepMind's Graph Nets).
    • MetaPath2Vec: Scalable representation learning for heterogeneous networks.
    • Graph Autoencoders (GAE/VGAE): Various implementations including Variational Graph Auto-Encoders and Adversarially Regularized versions.
    • SEAL: For link prediction tasks.
    • RENet: Recurrent Event Networks for temporal knowledge graphs.
    • GraphUNet: Graph U-Nets for hierarchical graph representation.
    • AttentiveFP: Graph attention mechanism optimized for molecular representation.
    • DeepGCN / GENConv: Layers designed for training deeper GCNs.
    • RECT: For network embedding with imbalanced labels.
    • LINKX: For large-scale learning on non-homophilous graphs.
    • RevGNN: Enables training extremely deep GNNs (up to 1000 layers).

    Knowledge Graph Embeddings (KGE)

    PyG includes several KGE models for multi-relational data:

    • TransE, ComplEx, DistMult, and RotatE.

    Model Explainability

    • GNNExplainer: An algorithm to generate explanations for GNN predictions by identifying important subgraphs.

  12. Generate explanations for Graph Neural Networks with torch_geometric.explain

    master

    The torch_geometric.explain package provides tools to explain GNN model predictions or to explain underlying phenomena within a dataset. It supports various explanation methods including:

    • GNNExplainer: Used for tasks like node classification and link prediction.
    • Captum-based explainers: Integrates with the Captum library for node classification and heterogeneous link prediction.
    • GraphMaskExplainer: A specialized explainer for node classification.
    • M-GNAN (MGNAN): An extension of GNAN (available in torch_geometric.contrib.nn.models.MGNAN) for graph classification that includes built-in node importance scores.