Torch Geometric

PyTorch Geometric is a community skill for building graph neural networks with PyTorch, covering message passing layers, graph convolutions, node and edge feature processing, graph pooling, and mini-batch handling for structured data learning.

What Is This?

Overview

PyTorch Geometric provides guidance on building graph neural networks using the PyG library for PyTorch. It covers message passing layers that propagate information between connected nodes through learnable aggregation functions, graph convolution operators including GCN, GAT, and GraphSAGE for neighborhood feature aggregation, node and edge feature processing that transforms attributes through graph-aware operations, graph pooling methods that create fixed-size representations from variable-size graphs for classification, and mini-batch handling that efficiently groups multiple graphs into single tensors for parallel training. The skill helps researchers and engineers apply deep learning to graph-structured data across domains including chemistry, biology, and social network analysis.

Who Should Use This

This skill serves machine learning researchers working with graph-structured datasets, data scientists analyzing social networks or molecular graphs, and engineers building recommendation systems on relational data. It is also relevant for scientists modeling physical simulations or knowledge graphs.

Why Use It?

Problems It Solves

Standard neural networks cannot directly process variable-size graph structures with arbitrary node connectivity. Implementing message passing from scratch requires careful sparse matrix operations and batching logic. Scaling graph neural networks to large graphs needs efficient neighbor sampling and mini-batching. Converting raw graph data into training-ready tensor formats requires specialized data loading pipelines.

Core Highlights

Message passing engine propagates features along graph edges. Convolution library provides standard GCN, GAT, and SAGE layers. Pooling operators create fixed-size graph-level representations. Batch loader groups multiple graphs for efficient training.

How to Use It?

Basic Usage

import torch
import torch.nn.functional as F
from torch_geometric.nn import (
    GCNConv,
    global_mean_pool
)
from torch_geometric.data import (
    Data
)

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

    def forward(
        self, x,
        edge_index, batch
    ):
        x = self.conv1(
            x, edge_index)
        x = F.relu(x)
        x = self.conv2(
            x, edge_index)
        x = global_mean_pool(
            x, batch)
        return self.lin(x)

Real-World Examples

from torch_geometric.nn import (
    GATConv
)
from torch_geometric.datasets\
    import Planetoid

dataset = Planetoid(
    root='data',
    name='Cora')
data = dataset[0]

class GAT(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.att1 = GATConv(
            dataset
                .num_features,
            8, heads=8)
        self.att2 = GATConv(
            64,
            dataset
                .num_classes,
            heads=1)

    def forward(
        self, x,
        edge_index
    ):
        x = F.elu(
            self.att1(
                x,
                edge_index))
        x = F.dropout(
            x, p=0.6,
            training=
                self.training)
        x = self.att2(
            x, edge_index)
        return F.log_softmax(
            x, dim=1)

Advanced Tips

Use NeighborLoader for sampling subgraphs during training on large graphs that do not fit in memory. Apply edge dropout as regularization to prevent overfitting on graph structure. Combine multiple convolution types with JumpingKnowledge connections for richer node representations. When working with heterogeneous graphs, use HeteroData objects and convert homogeneous models with the to_hetero utility to handle multiple node and edge types without rewriting layer logic.

When to Use It?

Use Cases

Classify molecules by predicting properties from atomic bond graphs. Detect communities in social networks using node embeddings. Build recommendation engines from user-item interaction graphs.

Related Topics

PyTorch, graph neural networks, GCN, GAT, message passing, node classification, and graph learning.

Important Notes

Requirements

PyTorch with torch-geometric and its sparse tensor dependencies installed for graph operations and message passing layers. Graph data formatted as edge index tensors and node feature matrices following the PyG Data object convention, with optional edge attributes and graph-level labels for supervised tasks. GPU with sufficient memory for training on large graphs with many nodes and dense feature vectors.

Usage Recommendations

Do: use DataLoader from torch_geometric for automatic mini-batching of multiple graphs into efficient batch tensors. Normalize node features before passing them to convolution layers, as unnormalized features can cause unstable training and slower convergence. Use built-in dataset classes for standard benchmarks to ensure reproducible comparisons.

Don't: stack too many convolution layers since deep GNNs suffer from over-smoothing where all node features converge. Ignore edge direction in directed graphs when it carries semantic meaning. Load entire large graphs into memory when neighbor sampling provides scalable training.

Limitations

Very deep graph networks suffer from over-smoothing that makes node representations indistinguishable. Large-scale graphs with millions of nodes require specialized sampling strategies that add complexity. Heterogeneous graphs with multiple node and edge types need specific layer variants and separate embedding spaces beyond standard homogeneous convolutions, increasing model complexity significantly.