pytorch_geometric README

The project's README from github.com/pyg-team/pytorch_geometric. As one markdown file: README.md.


PyPI Version PyPI Download Slack Contributing

Documentation | PyG 1.0 Paper | PyG 2.0 Paper | Colab Notebooks | External Resources | OGB Examples

PyG (PyTorch Geometric) is a library built upon PyTorch to easily write and train Graph Neural Networks (GNNs) for a wide range of applications related to structured data.

It consists of various methods for deep learning on graphs and other irregular structures, also known as geometric deep learning, from a variety of published papers. In addition, it consists of easy-to-use mini-batch loaders for operating on many small and single giant graphs, multi GPU-support, torch.compile support, DataPipe support, a large number of common benchmark datasets (based on simple interfaces to create your own), and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.

Click here to join our Slack community!


Library Highlights

Whether you are a machine learning researcher or first-time user of machine learning toolkits, here are some reasons to try out PyG for machine learning on graph-structured data.

Quick Tour for New Users

In this quick tour, we highlight the ease of creating and training a GNN model with only a few lines of code.

Train your own GNN model

In the first glimpse of PyG, we implement the training of a GNN for classifying papers in a citation graph. For this, we load the Cora dataset, and create a simple 2-layer GCN model using the pre-defined GCNConv:

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

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: Tensor, edge_index: Tensor) -> Tensor:
        # x: Node feature matrix of shape [num_nodes, in_channels]
        # edge_index: Graph connectivity matrix of shape [2, num_edges]
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

model = GCN(dataset.num_features, 16, dataset.num_classes)
We can now optimize the model in a training loop, similar to the standard PyTorch training procedure.
import torch.nn.functional as F

data = dataset[0]
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(200):
    pred = model(data.x, data.edge_index)
    loss = F.cross_entropy(pred[data.train_mask], data.y[data.train_mask])

    # Backpropagation
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

More information about evaluating final model performance can be found in the corresponding example.

Create your own GNN layer

In addition to the easy application of existing GNNs, PyG makes it simple to implement custom Graph Neural Networks (see here for the accompanying tutorial). For example, this is all it takes to implement the edge convolutional layer from Wang et al.:

$$x_i^{\prime} ~ = ~ \max_{j \in \mathcal{N}(i)} ~ \textrm{MLP}_{\theta} \left( [ ~ x_i, ~ x_j - x_i ~ ] \right)$$

import torch
from torch import Tensor
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")  # "Max" aggregation.
        self.mlp = Sequential(
            Linear(2 * in_channels, out_channels),
            ReLU(),
            Linear(out_channels, out_channels),
        )

    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
        # x: Node feature matrix of shape [num_nodes, in_channels]
        # edge_index: Graph connectivity matrix of shape [2, num_edges]
        return self.propagate(edge_index, x=x)  # shape [num_nodes, out_channels]

    def message(self, x_j: Tensor, x_i: Tensor) -> Tensor:
        # x_j: Source node features of shape [num_edges, in_channels]
        # x_i: Target node features of shape [num_edges, in_channels]
        edge_features = torch.cat([x_i, x_j - x_i], dim=-1)
        return self.mlp(edge_features)  # shape [num_edges, out_channels]

Architecture Overview

PyG provides a multi-layer framework that enables users to build Graph Neural Network solutions on both low and high levels. It comprises of the following components:

Implemented GNN Models

We list currently supported PyG models, layers and operators according to category:

GNN layers: All Graph Neural Network layers are implemented via the nn.MessagePassing interface. A GNN layer specifies how to perform message passing, i.e. by designing different message, aggregation and update functions as defined here. These GNN layers can be stacked together to create Graph Neural Network models.

Expand to see all implemented GNN layers...

Pooling layers: Graph pooling layers combine the vectorial representations of a set of nodes in a graph (or a subgraph) into a single vector representation that summarizes its properties of nodes. It is commonly applied to graph-level tasks, which require combining node features into a single graph representation.

Expand to see all implemented pooling layers...

GNN models: Our supported GNN models incorporate multiple message passing layers, and users can directly use these pre-defined models to make predictions on graphs. Unlike simple stacking of GNN layers, these models could involve pre-processing, additional learnable parameters, skip connections, graph coarsening, etc.

Expand to see all implemented GNN models...

GNN operators and utilities: PyG comes with a rich set of neural network operators that are commonly used in many GNN models. They follow an extensible design: It is easy to apply these operators and graph utilities to existing GNN layers and models to further enhance model performance.

Expand to see all implemented GNN operators and utilities...

Scalable GNNs: PyG supports the implementation of Graph Neural Networks that can scale to large-scale graphs. Such application is challenging since the entire graph, its associated features and the GNN parameters cannot fit into GPU memory. Many state-of-the-art scalability approaches tackle this challenge by sampling neighborhoods for mini-batch training, graph clustering and partitioning, or by using simplified GNN models. These approaches have been implemented in PyG, and can benefit from the above GNN layers, operators and models.

Expand to see all implemented scalable GNNs...

Installation

PyG is available for Python 3.10 to Python 3.14.

From PyG 2.3 onwards, you can install and use PyG without any external library required except for PyTorch. For this, simply run

pip install torch_geometric

Additional Libraries

If you want to utilize the full set of features from PyG, there exists several additional libraries you may want to install:

These packages come with their own CPU and GPU kernel implementations based on the PyTorch C++/CUDA/hip(ROCm) extension interface. For a basic usage of PyG, these dependencies are fully optional. We recommend to start with a minimal installation, and install additional dependencies once you start to actually need them.

For ease of installation of these extensions, we provide pip wheels for all major OS/PyTorch/CUDA combinations, see here.

PyTorch 2.12

To install the binaries for PyTorch 2.12, simply run

pip install pyg_lib torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.12.0+${CUDA}.html

where ${CUDA} should be replaced by either cpu, cu126, cu130, or cu132 depending on your PyTorch installation.

cpu cu126 cu128 cu130 cu132
Linux ✅ ✅ ✅ ✅
Windows ✅ ✅ ✅ ✅
macOS ✅

PyTorch 2.11

To install the binaries for PyTorch 2.11, simply run

pip install pyg_lib torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.11.0+${CUDA}.html

where ${CUDA} should be replaced by either cpu, cu126, cu128, or cu130 depending on your PyTorch installation.

cpu cu126 cu128 cu130
Linux ✅ ✅ ✅ ✅
Windows ✅ ✅ ✅ ✅
macOS ✅

PyTorch 2.10

To install the binaries for PyTorch 2.10, simply run

pip install pyg_lib torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.10.0+${CUDA}.html

where ${CUDA} should be replaced by either cpu, cu126, cu128, or cu130 depending on your PyTorch installation.

cpu cu126 cu128 cu130
Linux ✅ ✅ ✅ ✅
Windows ✅ ✅ ✅ ✅
macOS ✅

PyTorch 2.9

To install the binaries for PyTorch 2.9, simply run

pip install pyg_lib torch_scatter torch_sparse -f https://data.pyg.org/whl/torch-2.9.0+${CUDA}.html

where ${CUDA} should be replaced by either cpu, cu126, cu128, or cu130 depending on your PyTorch installation.

cpu cu126 cu128 cu130
Linux ✅ ✅ ✅ ✅
Windows ✅ ✅ ✅ ✅
macOS ✅

Note: Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1, PyTorch 1.13.0/1.13.1, PyTorch 2.0.0/2.0.1, PyTorch 2.1.0/2.1.1/2.1.2, PyTorch 2.2.0/2.2.1/2.2.2, PyTorch 2.3.0/2.3.1, PyTorch 2.4.0/2.4.1, PyTorch 2.5.0/2.5.1, PyTorch 2.6.0, PyTorch 2.7.0/2.7.1, and PyTorch 2.8.0 (following the same procedure). For older versions, you might need to explicitly specify the latest supported version number or install via pip install --no-index in order to prevent a manual installation from source. You can look up the latest supported version number here.

NVIDIA PyG Container

NVIDIA provides a PyG docker container for effortlessly training and deploying GPU accelerated GNNs with PyG, see here.

Nightly and Master

In case you want to experiment with the latest PyG features which are not fully released yet, either install the nightly version of PyG via

pip install pyg-nightly

or install PyG from master via

pip install git+https://github.com/pyg-team/pytorch_geometric.git

ROCm Wheels

The external pyg-rocm-build repository provides wheels and detailed instructions on how to install PyG for ROCm. If you have any questions about it, please open an issue here.

Cite

Please cite our PyG 1.0 and PyG 2.0 papers if you use this code in your own work:

@inproceedings{Fey/Lenssen/2019,
  title={Fast Graph Representation Learning with {PyTorch Geometric}},
  author={Fey, Matthias and Lenssen, Jan E.},
  booktitle={ICLR Workshop on Representation Learning on Graphs and Manifolds},
  year={2019},
}

@inproceedings{Fey/etal/2025,
  title={{PyG} 2.0: Scalable Learning on Real World Graphs},
  author={Fey, Matthias and Sunil, Jinu and Nitta, Akihiro and Puri, Rishi and Shah, Manan, and Stojanovi{\v{c}, Bla{\v{z} and Bendias, Ramona and Alexandria, Barghi and Kocijan, Vid and Zhang, Zecheng and He, Xinwei and Lenssen, Jan E. and Leskovec, Jure},
  booktitle={Temporal Graph Learning Workshop @ KDD},
  year={2025},
}

Feel free to email us if you wish your work to be listed in the external resources. If you notice anything unexpected, please open an issue and let us know. If you have any questions or are missing a specific feature, feel free to discuss them with us. We are motivated to constantly make PyG even better.

pytorch_geometric (pyg-team/pytorch_geometric)

Searching Database…

这是干什么的?PyG(PyTorch Geometric)是一个基于 PyTorch 的 Python 库,用于在结构化数据(图以及其他不规则结构)上编写和训练图神经网络(GNN),其中融合了大量已发表论文中的方法。
PyG 是一个独立的框架,还是在 PyTorch 之上使用的库?它是构建在 PyTorch 之上的一个库,并不是独立的框架:它保持了以张量为中心的 API,和原生 PyTorch 非常接近,所以如果你已经会写 PyTorch 代码,PyG 会让你一眼就觉得熟悉。
这是给研究人员用的,还是 PyTorch 新手也能上手?两者都适合:README 明确表示它既面向机器学习研究人员,也面向第一次接触机器学习工具包的用户,而且还有专门给新手准备的快速入门教程。
训练我的第一个模型需要写多少代码?非常少:快速入门教程只用 10 到 20 行代码就能在 Cora 引文图上训练一个 2 层 GCN,用的是现成的 GCNConv 层,训练循环也和原生 PyTorch 一模一样。
它只能处理普通图,还是也能处理其他数据?除了任意图之外,它还能处理三维网格和点云,并提供了专门的变换和层,比如来自 PointNet++ 系列工作的 PointNetConv 和 EdgeConv。
我能自己实现层吗,还是只能用内置的那些?完全可以自由扩展:在消息传递 API 里,一个层其实就是 message、aggregation 和 update 这几个函数——README 里就用几行代码演示了如何实现 EdgeConv 层。
它自带可以练手的数据集吗?有——它内置了大量常见的基准数据集,接口简单,还可以用来创建自己的数据集;配套的 Colab 笔记本让你直接在浏览器里运行示例。
What are similar projects?dgl, wuji-ego-mint, tapnet, lingbot-map, ParlAI, sam3, vegaml, UniMate
Creator
pyg-team
Size
24.4 MB
Stars
24.1K
Age
9y
Commits · 90 Days