Graph Neural Networks in PyTorch Without Extra Pain and Custom Workarounds

Most developers are used to working with clear, regular data structures. An image is a grid of pixels, and text easily fits into a sequence of tokens. Regular convolutional networks or transformers handle these perfectly.

Problems arise where data is fundamentally non-linear. Take banking transactions, connections between users in social networks, or chemical formulas of complex molecules. Graphs are at work everywhere. If you try to feed them to regular PyTorch, you'll have to manually struggle with sparse matrices, write complex loops for node aggregation, and keep track of indices.

This is exactly what PyTorch Geometric (abbreviated PyG) was created for. It's an extension for PyTorch that handles all the low-level math for you and provides an intuitive API for working with graph neural networks.

How PyG is Structured

If you already know how to write models in PyTorch, you'll get comfortable with PyG in a couple of hours. The framework follows the same principles: the same modules torch.nn, familiar training loops, and direct work with tensors.

All the magic of graph networks is built around the Message Passing concept. Each vertex in a graph collects information from its neighbors, combines it, and updates its own state.

PyG provides a base class MessagePassing that abstracts this logic. Inside, you only need to define three things:

• How a message is formed from a neighboring node
• How these messages are aggregated (sum, mean, max)
• How the node itself is updated

Here's what a simple graph convolution layer implementation looks like:

import torch
from torch.nn import Sequential, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr="max")
        self.mlp = Sequential(
            Linear(2 * in_channels, out_channels),
            ReLU(),
            Linear(out_channels, out_channels),
        )

    def forward(self, x, edge_index):
        # x задает фичи узлов, edge_index отвечает за связи между ними
        return self.propagate(edge_index, x=x)

    def message(self, x_j, x_i):
        # x_i — текущий узел, x_j — его сосед
        edge_features = torch.cat([x_i, x_j - x_i], dim=-1)
        return self.mlp(edge_features)

You don't need to write loops over all graph edges. PyG parallelizes this operation itself and executes it quickly thanks to compiled CUDA kernels.

Basic Example: Node Classification

Let's take the standard dataset Cora, which consists of scientific papers and citations between them. The task is to predict the category of a paper based on its text and connections to other publications.

A regular convolutional model GCN with data loading is assembled in literally two dozen lines:

import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv

# Загружаем датасет
dataset = Planetoid(root='.', name='Cora')

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

model = GCN(dataset.num_features, 16, dataset.num_classes)

The training loop looks completely standard for PyTorch enthusiasts: we pass the feature matrix x and the edge tensor edge_index, compute CrossEntropy, and call loss.backward().

The Main Problem of Graphs and How PyG Solves It

The main challenge of graph deep learning is scaling. Images in a batch can be easily split and sent to the GPU in parts. But in a graph, all nodes are interconnected. When a graph grows to millions of nodes and billions of edges, it no longer fits in the memory of even the most powerful GPUs.

PyG developers spent a lot of effort solving this problem. The library contains subgraph sampling mechanisms:

• NeighborLoader selects only random neighbors for each node in the batch, preventing memory from exploding
• ClusterGCN clusters a large graph into independent pieces and trains the network on them
• GraphSAINT samples random subgraphs while preserving their topology

Thanks to this, PyG can chew through giant graphs on regular GPUs without running out of memory.

What's Else Included Out of the Box

The repository contains an impressive collection of already implemented architectures from research papers. You'll find classic GCN, GAT, and GraphSAGE, as well as specialized models like SchNet and DimeNet for molecular analysis or PointNet for working with 3D point clouds.

Beyond algorithms, the library includes loaders for hundreds of standard datasets: from social networks to bioinformatics. This saves a lot of time on writing parsers.

Installation and Pitfalls

Since version 2.3, the base package installs extremely easily:

pip install torch_geometric

This is enough to get started. However, if you need maximum speed on huge graphs, you'll need to install additional libraries with C++/CUDA extensions: pyg-lib, torch-scatter, and torch-sparse.

This is where things sometimes get tricky. These binaries are tightly coupled to specific versions of PyTorch and CUDA. If you install an incompatible build, Python will flood you with C++ library import errors. Always check the compatibility matrix on the project website and install extensions with explicit wheel URLs.

Who This Framework Is For

The library is worth trying if your data doesn't fit well into tables or grids:

• Recommendation systems: the "user-item" relationship graph works better than regular matrix factorization.
• Anti-fraud and fintech: finding chains of suspicious transactions and hidden groups of fraudsters.
• Bioinformatics and chemistry: predicting molecular properties, drug discovery, and protein analysis.
• 3D data processing: LiDAR scans and point clouds are perfectly represented as graphs.

PyG has become the de facto standard in the industry for working with graphs. If you're planning to solve such problems, building custom architectures from scratch is definitely not worth it — PyG will save you a lot of time and resources.