github pyg-team/pytorch_geometric 2.5.0
PyG 2.5.0: Distributed training, graph tensor representation, RecSys support, native compilation
on GitHub

latest releases: 2.6.1, 2.6.0, 2.5.3...
7 months ago

We are excited to announce the release of PyG 2.5 🎉🎉🎉

  • Highlights
  • Breaking Changes
  • Deprecations
  • Features
  • Bugfixes
  • Changes
  • Full Changelog

PyG 2.5 is the culmination of work from 38 contributors who have worked on features and bug-fixes for a total of over 360 commits since torch-geometric==2.4.0.

Highlights

torch_geometric.distributed

We are thrilled to announce the first in-house distributed training solution for PyG via the torch_geometric.distributed sub-package. Developers and researchers can now take full advantage of distributed training on large-scale datasets which cannot be fully loaded in memory of one machine at the same time. This implementation doesn't require any additional packages to be installed on top of the default PyG stack.

Key Advantages

  • Balanced graph partitioning via METIS ensures minimal communication overhead when sampling subgraphs across compute nodes.
  • Utilizing DDP for model training in conjunction with RPC for remote sampling and feature fetching routines (with TCP/IP protocol and gloo communication backend) allows for data parallelism with distinct data partitions at each node.
  • The implementation via custom GraphStore and FeatureStore APIs provides a flexible and tailored interface for distributing large graph structure information and feature storage.
  • Distributed neighbor sampling is capable of sampling in both local and remote partitions through RPC communication channels. All advanced functionality of single-node sampling are also applicable for distributed training, e.g., heterogeneous sampling, link-level sampling, temporal sampling, etc.
  • Distributed data loaders offer a high-level abstraction for managing sampler processes, ensuring simplicity and seamless integration with standard PyG data loaders.

See here for the accompanying tutorial. In addition, we provide two distributed examples in examples/distributed/pyg to get started:

EdgeIndex Tensor Representation

torch-geometric==2.5.0 introduces the EdgeIndex class.

EdgeIndex is a torch.Tensor, that holds an edge_index representation of shape [2, num_edges]. Edges are given as pairwise source and destination node indices in sparse COO format. While EdgeIndex sub-classes a general torch.Tensor, it can hold additional (meta)data, i.e.:

  • sparse_size: The underlying sparse matrix size
  • sort_order: The sort order (if present), either by row or column
  • is_undirected: Whether edges are bidirectional.

Additionally, EdgeIndex caches data for fast CSR or CSC conversion in case its representation is sorted (i.e. its rowptr or colptr). Caches are filled based on demand (e.g., when calling EdgeIndex.sort_by()), or when explicitly requested via EdgeIndex.fill_cache_(), and are maintained and adjusted over its lifespan (e.g., when calling EdgeIndex.flip()). 

from torch_geometric import EdgeIndex

edge_index = EdgeIndex(
    [[0, 1, 1, 2],
     [1, 0, 2, 1]]
    sparse_size=(3, 3),
    sort_order='row',
    is_undirected=True,
    device='cpu',
)
>>> EdgeIndex([[0, 1, 1, 2],
...            [1, 0, 2, 1]])
assert edge_index.is_sorted_by_row
assert edge_index.is_undirected

# Flipping order:
edge_index = edge_index.flip(0)
>>> EdgeIndex([[1, 0, 2, 1],
...            [0, 1, 1, 2]])
assert edge_index.is_sorted_by_col
assert edge_index.is_undirected

# Filtering:
mask = torch.tensor([True, True, True, False])
edge_index = edge_index[:, mask]
>>> EdgeIndex([[1, 0, 2],
...            [0, 1, 1]])
assert edge_index.is_sorted_by_col
assert not edge_index.is_undirected

# Sparse-Dense Matrix Multiplication:
out = edge_index.flip(0) @ torch.randn(3, 16)
assert out.size() == (3, 16)

EdgeIndex is implemented through extending torch.Tensor via the __torch_function__ interface (see here for the highly recommended tutorial).

EdgeIndex ensures for optimal computation in GNN message passing schemes, while preserving the ease-of-use of regular COO-based PyG workflows. EdgeIndex will fully deprecate the usage of SparseTensor from torch-sparse in later releases, leaving us with just a single source of truth for representing graph structure information in PyG.

RecSys Support

Previously, all/most of our link prediction models were trained and evaluated using binary classification metrics. However, this usually requires that we have a set of candidates in advance, from which we can then infer the existence of links. This is not necessarily practical, since in most cases, we want to find the top-k most likely links from the full set of O(N^2) pairs.

torch-geometric==2.5.0 brings full support for using GNNs as a recommender system (#8452), including support for 

mips = MIPSKNNIndex(dst_emb)

for src_batch in src_loader:
    src_emb = model(src_batch.x_dict, src_batch.edge_index_dict)
    _, pred_index_mat = mips.search(src_emb, k)
    
    for metric in retrieval_metrics:
         metric.update(pred_index_mat, edge_label_index)
         
for metric in retrieval_metrics:
     metric.compute()

See here for the accompanying example.

PyTorch 2.2 Support

PyG 2.5 is fully compatible with PyTorch 2.2 (#8857), and supports the following combinations:

PyTorch 2.2 cpu cu118 cu121
Linux ✅ ✅ ✅
macOS ✅
Windows ✅ ✅ ✅

You can still install PyG 2.5 with an older PyTorch release up to PyTorch 1.12 in case you are not eager to update your PyTorch version.

Native torch.compile(...) and TorchScript Support

torch-geometric==2.5.0 introduces a full re-implementation of the MessagePassing interface, which makes it natively applicable to both torch.compile and TorchScript. As such, torch_geometric.compile is now fully deprecated in favor of torch.compile 

- model = torch_geometric.compile(model)
+ model = torch.compile(model)

and MessagePassing.jittable() is now a no-op: 

- conv = torch.jit.script(conv.jittable())
+ model = torch.jit.script(conv)

In addition, torch.compile usage has been fixed to not require disabling of extension packages such as torch-scatter or torch-sparse.

New Tutorials, Examples, Models and Improvements

Breaking Changes

  • GATConv now initializes modules differently depending on whether their input is bipartite or non-bipartite (#8397). This will lead to issues when loading model state for GATConv layers trained on earlier PyG versions.

Deprecations

Features

Package-wide Improvements

Temporal Graph Support

torch_geometric.datasets

torch_geometric.nn

torch_geometric.metrics

torch_geometric.explain

torch_geometric.transforms

Other Improvements

  • Added support for returning multi graphs in utils.to_networkx (#8575)
  • Added noise scheduler utilities utils.noise_scheduler.{get_smld_sigma_schedule,get_diffusion_beta_schedule} for diffusion-based graph generative models (#8347)
  • Added a relabel node functionality to utils.dropout_node via relabel_nodes: bool argument (#8524)
  • Added support for weighted utils.cross_entropy.sparse_cross_entropy (#8340)
  • Added support for profiling on XPU device via profile.profileit("xpu") (#8532)
  • Added METIS partitioning with CSC/CSR format selection in ClusterData (#8438)

Bugfixes

Changes

New Contributors

Full Changelog: 2.4.0...2.5.0

Check out latest releases or
releases around pyg-team/pytorch_geometric 2.5.0

Don't miss a new pytorch_geometric release

NewReleases is sending notifications on new releases.

Get notifications