Graph neural networks occupy a peculiar spot in deep learning research: everyone agrees they matter, yet the mechanics of a GNN layer feel far more opaque than those of a convolution or a transformer block. Images come with a fixed grid and a battle-tested toolkit; graphs bring arbitrary topology, ragged neighborhoods, and batches that are not tensors at all. PyTorch Geometric — usually abbreviated PyG — is the library that made this manageable on PyTorch, and it is now the de facto standard for GNN work in the PyTorch ecosystem, powering everything from citation-network tutorials to drug-discovery and recommender systems in production.

The library deserves a source tour for the same reason PyTorch itself does: its most valuable ideas are not in the README but in the way a handful of abstractions interlock. PyG is built around four load-bearing pillars. First, a Data object that represents one graph as plain tensors plus an edge_index matrix. Second, a data loading flow that can either stack whole graphs into mega-batches or sample subgraphs out of graphs far too large to fit in memory. Third, a message-passing engine that reduces every published GNN layer — GCN, GraphSAGE, GAT and hundreds more — to a small set of overridable functions. Fourth, scatter-based aggregation kernels that do the heavy lifting whenever many edges must be reduced onto the same node.

Understanding how these pillars connect explains most of what makes PyG feel the way it does in daily use: why handing a DataLoader a list of Data objects yields a Batch, why sampling loaders return ever-growing node sets per hop, and why a custom GNN layer is often under forty lines of code. This post walks the actual source tree of the repository, with diagrams and line-level links, so you can use it as a map the next time a training loop misbehaves or a paper’s equation needs to become a layer.

The five load-bearing regions of PyG: the Data layer, the loading and sampling flow, the message-passing family, scatter-based aggregation and pooling, and the ready-made GNN models built on top of all four.

Reading the overview from left to right:

• Everything starts at Dataset, L30, the abstract base that defines the download/process/len/getitem protocol every PyG dataset implements. The in-memory flavor at InMemoryDataset, L32 caches a whole corpus as a single collated tensor store.
• One graph is one Data, L472 instance: x for node features, edge_index for connectivity in COO format, plus arbitrary extra tensors.
• The classic DataLoader, L52 subclasses torch.utils.data.DataLoader and swaps the collation function so that a list of graphs becomes one parallel Batch object.
• When a single graph is too big to batch, NeighborLoader, L10 delegates to a sampling engine descended from BaseSampler, L932 and returns subgraph-shaped Data objects instead.
• Every convolution in the library descends from MessagePassing, L39. Concrete layers such as GCNConv, L116, SAGEConv, L13 and GATConv, L35 differ only in how they compute and reduce messages.
• The reduction itself flows through Aggregation, L10, while graph-level classification tasks use readouts like global_mean_pool, L37.
• Finally, batteries-included classifiers like the GCN model, L389 stack conv layers behind a uniform training interface, so a baseline is one constructor call away.

From Install to Your First Model

PyG installs like any other PyTorch extension: pip install torch_geometric works out of the box with recent PyTorch versions, and the optional accelerators torch_scatter, torch_sparse, torch_cluster and torch_spline_conv only buy you speed and a few extra algorithms — the pure-PyTorch core works without a compiler.

The canonical hello world trains a two-layer GCN on Cora, a citation graph with about 2,700 paper nodes and seven classes. The dataset side is one line, because Planetoid, L11 subclasses InMemoryDataset and downloads, processes and caches everything on first use:

from torch_geometric.datasets import Planetoid
dataset = Planetoid(root="/tmp/Cora", name="Cora")
data = dataset[0]           # one big Data object, train/val/test masks included

The model side shows why the template BasicGNN classes exist:

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCN

model = GCN(dataset.num_features, 16, dataset.num_classes, num_layers=2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

for epoch in range(200):
    out = model(data.x, data.edge_index)   # logits for every node
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    optimizer.zero_grad(); loss.backward(); optimizer.step()

That loop is full-batch training: the whole graph goes through the network each step, and the GCN model from basic_gnn.py handles the conv stack, dropout and classification head for you. If you would rather assemble layers yourself, the explicit equivalent uses GCNConv, L116 directly — GCNConv(in, out) twice with F.relu between them — and behaves identically. One small but telling detail: calling seed_everything, L7 before the loop reseeds Python, NumPy and PyTorch in one shot — handy, because Cora runs are annoyingly unstable otherwise.

The Data Layer: Graphs as Tensors with Topology

The entire input side of PyG rests on Data, L472. Its central convention is edge_index: an int64 tensor of shape [2, num_edges] holding source and destination node indices. Nothing else about a graph is mandatory: node labels are just another tensor, and heterogeneous graphs extend the same idea with string-qualified node and edge types — the typing module formalizes those as NodeType and EdgeType aliases. This “dictionary of tensors” design is why PyG composes so cleanly with PyTorch autograd: gradients flow through x and edge weights like through any tensor operation.

Datasets wrap collections of Data objects. The split between Dataset, L30 and InMemoryDataset, L32 is a deliberate memory trade: the former keeps a Python list and loads items lazily from disk, while the latter calls collate() once and stores the entire corpus as concatenated tensors with index pointers — precisely how Planetoid and most benchmark suites ship. When you then iterate a loader, collate, L37 performs the trick that makes batching graphs possible at all: node and edge tensors are concatenated along dimension zero, and a batch vector records which graph each node came from. The result, Batch, L57, is itself a Data subclass, so downstream code never needs a special case — a batch of 32 graphs is just one disconnected graph with 32 components.

Transforms plug into this layer the same way torchvision transforms plug into images. Anything implementing BaseTransform, L6 can be attached to a dataset, and Compose, L7 chains several together; NormalizeFeatures, L9 is the one every Cora tutorial uses. Because transforms run inside the dataset, preprocessing happens once at load time instead of per epoch.

Inside the Message-Passing Chain

If you read only one file in PyG, make it message_passing.py. The insight at MessagePassing, L39 is that virtually every GNN layer from the literature fits the same three-step template: compute a message per edge from the endpoint features, aggregate messages per destination node with a permutation-invariant reduction, and update node embeddings with the aggregated result. The base class turns that template into a dispatch loop — propagate, L421 gathers endpoint features, invokes message, L565, then routes the result into aggregate, L577 or, when a fused scatter path is available, straight into message_and_aggregate, L598.

The concrete layers show how little code the template actually needs. GCNConv implements the Kipf-Welling normalization by dividing messages by the square root of endpoint degrees — which it obtains from degree, L9 — and caches the normalized edge weights when the topology is static. SAGEConv concatenates the destination node’s own features with the mean of its neighbors’ messages, exactly the GraphSAGE recipe. GATConv scores each edge with a learned attention coefficient, normalizes those scores with the per-neighbor softmax from softmax, L12, and uses them to weight the aggregation — meaning attention itself is just another differentiable message weighting. All three funnel their reductions through the same scatter kernel at scatter, L14, which dispatches to whichever compiled backend is installed.

The detail view: how Data objects flow through collate and the loaders, how the sampler machinery sits behind NeighborLoader, and how every conv layer leans on the scatter, softmax and degree kernels before handing pooled readouts to the model templates.

The detail diagram above is the map worth keeping open while debugging. It makes two things visually obvious. First, the data loading flow is a funnel — Dataset transforms feed Data, collate feeds Batch, and both the vanilla DataLoader and the sampling loaders sit in front of the conv layers, so shape errors almost always originate left of the model. Second, aggregation is a genuinely separate subsystem: MessagePassing has a default reduction, but any layer accepts a custom Aggregation, L10 module, and simple ones like MeanAggregation, L25 and SumAggregation, L12 are two-liners worth copying when you experiment.

Loading and Sampling at Scale

Full-batch training dies the moment your graph has millions of nodes, so PyG’s answer is mini-batching by neighborhood. The vanilla DataLoader, L52 batches multiple graphs; NeighborLoader, L10 batches parts of one graph. Under the hood it asks the sampler hierarchy — BaseSampler, L932 with the concrete NeighborSampler, L40 — to walk num_neighbors random in-edges per node per hop, then materializes the induced subgraph as a fresh Data object. What you receive in the training loop is the same tensor interface as before, just with progressively larger node sets per layer — the subgraph’s “frontier”.

This shape is worth internalizing because it explains GraphSAGE’s original claim to fame: with two hops of fifteen-and-ten neighbors, a node’s representation depends on a bounded, size-independent slice of the graph, so memory per batch stays constant no matter how big the full adjacency matrix is. The cost is the classic trade-off — sampled training is unbiased only in expectation, and deep stacks with many hops suffer the “neighborhood explosion” where a single seed node drags in thousands of ancestors. Practical guidance falls out of the source structure: keep hops shallow, shrink num_neighbors as depth grows, and reach for the sampling loaders only when the full graph genuinely will not fit — the HuggingFace Datasets library plays a complementary role when your node features, not the graph itself, are the memory hog.

Aggregation and the Scatter Kernel

Every reduce-with-respect-to-index operation in PyG bottoms out at scatter, L14: given values and an index vector, add/mean/max them into slots. This is the GNN analog of matrix multiplication for transformers — the primitive whose efficiency decides the whole library’s throughput. Graph-level tasks add one more consumer: global_mean_pool, L37 runs the same scatter over the batch vector to produce one embedding per graph, which is what turns node-level architectures into graph classifiers, much like the pooled features from scikit-learn reducers.

The Aggregation base class turns that primitive into a first-class extension point. Subclasses implement one reduce call; utilities for multi-head combinations, sorting-based pooling and attention-weighted aggregation all build on it. When a paper proposes a new reduction — power means, soft maximum, anything — the PyG-native way to try it is a small Aggregation subclass rather than a forked conv layer, and the message-passing chain will pick it up through the standard aggr argument.

Honest Limits

Sampling is unbiased in theory and subtly biased in practice. Neighbor sampling’s variance shows up as degraded accuracy on tasks sensitive to high-degree nodes, and several papers report noticeable gaps between full-batch and sampled training on the same model; the library gives you the knobs (num_neighbors, disjoint vs. non-disjoint modes) but not a magic equalizer. The neighborhood explosion also inverts on very dense graphs, where “sample ten neighbors per hop” still yields enormous frontiers.

Heterogeneous graphs work but demand care. The HeteroData variant is genuinely well designed, yet real-world heterogeneous workflows involve a combinatorial spread of edge types, per-type sampling budgets and the notorious to_heterogeneous/to_homogeneous conversions whose failure modes are shape errors deep inside the message-passing chain. Expect to spend a day on schema design before the first epoch runs.

Finally, the dependency surface has sharp edges worth knowing before deployment. The optional acceleration libraries pin specific PyTorch and CUDA combinations, and performance-oriented code paths (fused scatter in particular) differ measurably between setups. Convolutions also evolve faster than tutorials: APIs like the pooling helpers moved across modules over the library’s history — the nn/glob.py path you will still find in old Stack Overflow answers is now a deprecation shim over the real implementation in nn/pool/glob.py. Pin your version, read the changelog on upgrade, and this library’s source will repay the read many times over.

Enjoyed this post? Never miss out on future posts by following us