zoom_in Click to enlarge

PyTorch Geometric

▲ ▼

language

description PyTorch Geometric Overview

For data structured as graphs (social networks, molecular structures, knowledge graphs), PyTorch Geometric (PyG) is the specialized tool. It extends PyTorch to handle graph convolutions and message passing efficiently. It is essential for any domain where relationships between entities are more important than the entities themselves, providing specialized layers for graph data science.

help PyTorch Geometric FAQ

How is PyTorch Geometric different from standard PyTorch for graph data?

PyTorch Geometric (PyG) introduces a specialized `Data` object that efficiently stores node features, edge indices, and labels in sparse COO format, which standard PyTorch tensors cannot handle compactly. It also provides optimized implementations of message-passing operators, so operations like graph convolutions run without the user manually writing scatter-gather logic.

Which graph neural network layers come pre-implemented in PyG?

PyG includes ready-to-use layers for GCN (Kipf & Welling), GraphSAGE, GAT (Graph Attention Networks), GIN (Graph Isomorphism Networks), and ChebConv, among many others. These layers follow the standard `torch.nn.Module` interface, so you can stack them into custom architectures exactly as you would with standard PyTorch layers.

Can PyG train on very large graphs that don't fit in GPU memory?

Yes, PyG provides `NeighborLoader` and `ClusterLoader` implementations that perform neighborhood sampling for mini-batch training on graphs with millions of nodes. This approach, commonly benchmarked on datasets like ogbn-products, samples a local subgraph around target nodes rather than loading the entire adjacency matrix into memory.

How do I load benchmark datasets like Cora or PubMed in PyTorch Geometric?

PyG includes the `Planetoid` dataset class accessible via `torch_geometric.datasets.Planetoid`, which automatically downloads the Cora, CiteSeer, and PubMed citation network datasets. These are the most commonly used benchmark graphs in GNN research and load as ready-to-train `Data` objects with node features, edge lists, and train/validation/test masks.

swap_horiz

Looking for PyTorch Geometric alternatives? Compare top competitors ranked & scored

arrow_forward