Pyg 2.3.0: PyTorch 2.0 support, native sparse tensor support, explainability and accelerations
We are thrilled to announce the release of PyG 2.3 ๐
Unclaimed project
Are you a maintainer of pytorch_geometric? Claim this project to take control of your public changelog and roadmap.
Changelog
Graph Neural Network Library for PyTorch
Last updated about 1 month ago
We are thrilled to announce the release of PyG 2.3 ๐
AutoGPT is the vision of accessible AI for everyone, to use and to build on. Our mission is to provide the tools, so that you can focus on what matters.
Stable Diffusion web UI
๐ค Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.
A feature-rich command-line audio/video downloader
PyG 2.3 is the culmination of work from 59 contributors who have worked on features and bug-fixes for a total of over 470 commits since torch-geometric==2.2.0.
PyG 2.3 is fully compatible with the next generation release of PyTorch, bringing many new innovations and features such as torch.compile() and Python 3.11 support to PyG out-of-the-box. In particular, many PyG models and functions are speeded up significantly using torch.compile() in torch >= 2.0.0.
We have prepared a full tutorial and a set of examples to get you going with torch.compile() immediately:
import torch_geometric
from torch_geometric.nn import GraphSAGE
model = GraphSAGE(in_channels, hidden_channels, num_layers, out_channels)
model = model.to(device)
model = torch_geometric.compile(model)
Overall, we observed runtime improvements of nearly up to 300%:
Model | Mode | Forward | Backward | Total | Speedup
-- | -- | -- | -- | -- | --
GCN | Eager | 2.6396s | 2.1697s | 4.8093s | ย
GCN | Compiled | 1.1082s | 0.5896s | 1.6978s | 2.83x
GraphSAGE | Eager | 1.6023s | 1.6428s | 3.2451s | ย
GraphSAGE | Compiled | 0.7033s | 0.7465s | 1.4498s | 2.24x
GIN | Eager | 1.6701s | 1.6990s | 3.3690s | ย
GIN | Compiled | 0.7320s | 0.7407s | 1.4727s | 2.29x
Please note that torch.compile() within PyG is in beta mode and under active development. For example, currently torch.compile(model, dynamic=True) does not yet work seamlessly, but fixes are on its way. We are very eager to improve its support across the whole PyG code base, so do not hesitate to reach out if you notice anything unexpected.
With the recent upstreams of torch-scatter and torch-sparse to native PyTorch, we are happy to announce that any installation of the extension packages torch-scatter, torch-sparse, torch-cluster and torch-spline-conv is now fully optional.
All it takes to install PyG is now encapsulated into a single command
pip install torch-geometric
and finally resolves a lot of previous installation issues.
Extension packages are still picked up for the following use-cases (if installed):
pyg-lib: Heterogeneous GNN operators and graph sampling routines like NeighborLoadertorch-scatter: Accelerated "min" and "max" reductionstorch-sparse: SparseTensor supporttorch-cluster: Graph clustering routines like knn or radiustorch-spline-conv: SplineConv supportWe recommend to start with a minimal installation, and only install additional dependencies once you actually get notified about them being missing during PyG usage.
With the recent additions of torch.sparse_csr_tensor and torch.sparse_csc_tensor classes and accelerated sparse matrix multiplication routines to PyTorch, we finally enable MessagePassing on pure PyTorch sparse tensors as well. In particular, you can now use torch.sparse_csr_tensor and torch.sparse_csc_tensor as a drop-in replacement for torch_sparse.SparseTensor:
from torch_geometric.nn import GCN
import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
transform = T.ToSparseTensor(layout=torch.sparse_csr)
dataset = Planetoid("Planetoid", name="Cora", transform=transform)
model = GCN(in_channels, hidden_channels, num_layers=2)
model = model(data.x, data.adj_t)
Nearly all of the native PyG layers have been tested to work seamlessly with native PyTorch sparse tensors (#5906, #5944, #6003, #6033, #6514, #6532, #6748, #6847, #6868, #6874, #6897, #6930, #6932, #6936, #6937, #6939, #6947, #6950, #6951, #6957).
In PyG 2.2 we introduced the torch_geometric.explain package that provides a flexible interface to generate and visualize GNN explanations using various algorithms. We are happy to add the following key improvements on this front:
HeteroExplanationvisualize_feature_importance and to visualize_graph explanationsCaptumExplainer, PGExplainer, AttentionExplainer, PGMExplainer, and GraphMaskExplainerUsing the new explainer interface is as simple as:
explainer = Explainer(
model=model,
algorithm=CaptumExplainer('IntegratedGradients'),
explanation_type='model',
model_config=dict(
mode='multiclass_classification',
task_level='node',
return_type='log_probs',
),
node_mask_type='attributes',
edge_mask_type='object',
)
explanation = explainer(data.x, data.edge_index)
Read more about torch_geometric.explain in our newly added tutorial and example scripts. We also added a blog post that describes the new interface and functionality in depth.
Together with Intel and NVIDIA, we are excited about new PyG accelerations:
[Experimental] Support for native cugraph-based GNN layers for commonly used layers such as CuGraphSAGEConv, CuGraphGATConv, and CuGraphRGCNConv (#6278, #6388, #6412):
RGCN with neighbor sampling:
Dataset | CuGraphRGCNConv (ms) | FastRGCNConv (ms) | RGCNConv (ms)
-- | -- | -- | --
AIFB | 7,2 | 13,4 | 70
BGS | 7,1 | 8,8 | 146,9
MUTAG | 8,3 | 21,8 | 47,6
AM | 17,5 | 51 | 330,1
Full-batch GAT:
Dataset | CuGraphGATConv (ms) | GATConv (ms)
-- | -- | --
Cora | 7 | 8,7
Citeseer | 7 | 9
Pubmed | 8,2 | 11,4
GraphSAGE with neighbor sampling:
Dataset | CuGraphSAGEConv (ms) | SAGEConv (ms)
-- | -- | --
ogbn-products | 2591,8 | 3040,3
A fast alternative to HGTConv via FastHGTConv that utilizes pyg-lib integration for improved runtimes. Overall, FastHGTConv achieves a speed-up of approximately 300% compared to the original implementation (#6178).
A fast implementation of HeteroDictLinear that utilizes pyg-lib integration for improved runtimes (#6178).
GNN inference and training optimizations on CPU within native PyTorch 2.0. Optimizations include:
scatter_reduce: performance hotspot in message passing when edge_index is stored in Coordinate format (COO).gather: backward of scatter_reduce, specially tuned for the GNN compute when the index is an expanded tensor.torch.sparse.mm with reduce flag: performance hotspot in message passing when the edge_index is stored in Compressed Sparse Row (CSR). Supported reduce flags are "sum", "mean", "amax" and "amin".
On OGB benchmarks, a 1.12x - 4.07x performance speedup is measured (PyTorch 1.13.1 vs PyTorch 2.0) for single node inference and training.Introduction of index_sort via pyg-lib>=0.2.0, which implements a (way) faster alternative to sorting one-dimensional indices compared to torch.sort (#6554). Overall, this achieves speed-ups in dataset loading times by up to 600%.
Introduction of AffinityMixin to accelerate PyG workflows on CPU. CPU affinity can be enabled via the AffinityMixin.enable_cpu_affinity() method for num_workers > 0 data loading use-cases, and will guarantee that a separate core is assigned to each worker at initialization. Over all benchmarked model/dataset samples, the average training time is decreased by up to 1.85x. We added an in-depth tutorial on how to speed-up your PyG workflows on CPU.
The documentation has undergone a revision of design and structure, making it faster to load and easier to navigate. Take a look at its new design here.
We had our third community sprint in the last two weeks of January. The goal was to improve code coverage by writing more thorough tests. Thanks to the efforts of many contributors, the total code coverage went from ~85% to ~92% (#6528, #6523, #6538, #6555, #6558, #6568, #6573, #6578, #6597, #6600, #6618, #6619, #6621, #6623, #6637, #6638, #6640, #6645, #6648, #6647, #6653, #6657, #6662, #6664, #6667, #6668, #6669, #6670, #6671, #6673, #6675, #6676, #6677, #6678, #6681, #6683, #6703, #6720, #6735, #6736, #6763, #6781, #6797, #6799, #6824, #6858)
NeighborLoader will now also sample nodes with an equal timestamp to the seed time. Changed from sampling only nodes with a smaller timestamp (requires pyg-lib>=0.2.0) (#6517)GraphMultisetTransformer such that GNN execution is no longer performed inside its module (#634))Explanation.node_mask and Explanation.node_feat_mask into a single attribute in Explainer (#6267)ExplainerConfig arguments to the Explainer class (#6176)torch_geometric.data.lightning (#6140)target_index argument in the Explainer interface (#6270)Aggregation.set_validate_args option (#6175)__dunder__ names in MessagePassing (#6999)datasets.BAShapes is now deprecated. Use the BAGraph graph generator to generate Barabasi-Albert graphs instead (#6072)DenseGATConv layer (https://github.com/pyg-team/pytorch_geometric/pull/6928)DistMult KGE model (https://github.com/pyg-team/pytorch_geometric/pull/6958)ComplEx KGE model (#6898)TransE KGE model (#6314)HeteroLayerNorm and HeteroBatchNorm layers (#6838)TemporalEncoding module (#6785)SimpleConv to perform non-trainable propagation (#6718)torch.jit examples for example/film.py and example/gcn.py (#6602)AntiSymmetricConv layer (#6577)PyGModelHubMixin for Huggingface model hub integration (#5930, #6591)PGMExplainer (#6149, #6588, #6589)ToHeteroLinear and ToHeteroMessagePassing modules to accelerate to_hetero functionality (#5992, #6456)GraphMaskExplainer (#6284)GRBCDAttack and PRBCDAttack adversarial attack models (#5972)CaptumExplainer (#6383, #6387, #6433, #6487)GNNFF model (#5866)MLPAggregation, SetTransformerAggregation, GRUAggregation, and DeepSetsAggregation as adaptive readout functions (#6301, #6336, #6338)
(https://github.com/pyg-team/pytorch_geometric/pull/6331), )GPSConv Graph Transformer layer (#6326, #6327)PGExplainer (#6204)AttentionExplainer (#6279)PointGNNConv layer (#6194)AirfRANS dataset (#6287)HeterophilousGraphDataset suite (#6846)MD17 dataset (#6734)BAMultiShapesDataset (#6541)Taobao dataset and a corresponding example (#6144)FB15k_237 dataset (#3204)BA2MotifDataset explainer dataset (#6257)CycleMotif motif generator to generate n-node cycle shaped motifs (#6256)InfectionDataset to evaluate explanations (#6222)CustomMotif motif generator (#6179)ERGraph graph generator to generate Ergos-Renyi (ER) graphs (#6073)ExplainerDataset to evaluate explanation methods (#6104)NeighborLoader to return the number of sampled nodes and edges per hop, and added corresponding trim_to_layer functionality for more efficient NeighborLoader use-cases (#6661, #6834)ZipLoader to execute multiple NodeLoader or LinkLoader instances (#6829)seed_time attribute to temporal NodeLoader outputs in case input_time is given (#6196)utils.one_hot implementation (https://github.com/pyg-team/pytorch_geometric/pull/7005)utils.softmax implementation (#6113, #6155, #6805)topk implementation for graph pooling on large graphs (#6123)utils.select and utils.narrow functionality to support filtering of both tensors and lists (#6162)normalization customization in get_mesh_laplacian (#6790)spmm functionality via CSR format (#6699, #6759)RECT_L model (#6727)Node2Vec model (#6726)utils.to_edge_index to convert sparse tensors to edge indices and edge attributes (#6728)LINKX model (#6712)dropout option to GraphMultisetTransformer (#6484)LightningNodeData and LightningLinkData (#6450, #6456)num_neighbors in NeighborSampler after instantiation (#6446)HeteroData mini-batch class in remote backends (#6377)ChebConv within GNNExplainer (https://github.com/pyg-team/pytorch_geometric/pull/6778)Dataset.to_datapipe functionality for converting PyG datasets into a PyTorch DataPipe(#6141)to_nested_tensor and from_nested_tensor functionality (#6329, #6330, [#6331]networkit conversion utilities (#6321)Data.update() and HeteroData.update functionality (#6313)HeteroData.set_value_dict functionality (https://github.com/pyg-team/pytorch_geometric/pull/6961, https://github.com/pyg-team/pytorch_geometric/pull/6974)fidelity explainability metric (#6116, #6510)characterization_score and fidelity_curve_auc explainer metrics (#6188)LinkNeighborLoader (#6264)get_embeddings function (#6201)Explanation.visualize_feature_importance to support node feature importance visualizations (#6094)HeteroExplanation (#6091, #6218)summary method for PyG/PyTorch models (#5859, #6161)input_time option to LightningNodeData and transform_sampler_output to NodeLoader and LinkLoader (#6187)Data.edge_subgraph and HeteroData.edge_subgraph functionalities (#6193)Data.subgraph() and HeteroData.subgraph() for bipartite graphs (#6613, #6654)PNAConv and DegreeScalerAggregation to correctly incorporate degree statistics of isolated nodes (#6609)Data.to_heterogeneous filtered attributes in the wrong dimension (#6522)to_hetero when using an uninitialized submodule without implementing reset_parameters (#6863)get_mesh_laplacian (#6790)GNNExplainer on link prediction tasks (#6787)ImbalancedSampler when operating on a sliced InMemoryDataset (#6374)transforms.GDC to not crash on graphs with isolated nodes (#6242)transforms.RemoveIsolatedNodes (#6308)DimeNet that causes an output dimension mismatch (#6305)Data.to_heterogeneous when used with an empty edge_index (#6304)HeteroLinear for un-sorted type vectors (#6198)utils.one_hot implementation (#7005)HeteroDictLinear and an optimized FastHGTConv module (#6178, #6998)DenseGATConv module (#6928)trim_to_layer utility function for more efficient NeighborLoader use-cases (#6661)DistMult KGE model (#6958)HeteroData.set_value_dict functionality (#6961, #6974)ComplEx KGE model (#6898)HeteroLayerNorm and HeteroBatchNorm layers (#6838)HeterophilousGraphDataset suite (#6846)NeighborLoader to return number of sampled nodes and edges per hop (#6834)ZipLoader to execute multiple NodeLoader or LinkLoader instances (#6829)utils.select and utils.narrow functionality to support filtering of both tensors and lists (#6162)normalization customization in get_mesh_laplacian (#6790)TemporalEncoding module (#6785)spmm_reduce functionality via CSR format (#6699, #6759)MD17 dataset (#6734)RECT_L model (#6727)Node2Vec model (#6726)utils.to_edge_index to convert sparse tensors to edge indices and edge attributes (#6728)PolBlogs dataset (#6714)SimpleConv to perform non-trainable propagation (#6718)RemoveDuplicatedEdges transform (#6709)LINKX model (#6712)torch.jit examples for example/film.py and example/gcn.py(#6602)Pad transform (#5940, #6697, #6731, #6758)cat aggregation type to the HeteroConv class so that features can be concatenated during grouping (#6634)torch.compile support and benchmark study (#6610, #6952, #6953, #6980, #6983, #6984, #6985, #6986, #6989, #7002)AntiSymmetricConv layer (#6577)nn.conv.cugraph via cugraph-ops (#6278, #6388, #6412)index_sort function from pyg-lib for faster sorting (#6554)EquilibriumAggregration (#6560)dense_to_sparse() (#6546)BAMultiShapesDataset (#6541)n_id and e_id attributes to mini-batches produced by NodeLoader and LinkLoader (#6524)PGMExplainer to torch_geometric.contrib (#6149, #6588, #6589)NumNeighbors helper class for specifying the number of neighbors when sampling (#6501, #6505, #6690)is_node_attr() and is_edge_attr() calls (#6492)ToHeteroLinear and ToHeteroMessagePassing modules to accelerate to_hetero functionality (#5992, #6456)GraphMaskExplainer (#6284)GRBCD and PRBCD adversarial attack models (#5972)dropout option to SetTransformer and GraphMultisetTransformer (#6484)LightningNodeData and LightningLinkData (#6450, #6456)num_neighbors in NeighborSampler after instantiation (#6446)Taobao dataset and a corresponding example for it (#6144)pyproject.toml (#6431)torch_geometric.contrib sub-package (#6422)pyright type checker support (#6415)CaptumExplainer (#6383, #6387, #6433, #6487, #6966)HeteroData mini-batch class in remote backends (#6377)GNNFF model (#5866)MLPAggregation, SetTransformerAggregation, GRUAggregation, and DeepSetsAggregation as adaptive readout functions (#6301, #6336, #6338)Dataset.to_datapipe for converting PyG datasets into a torchdata DataPipe(#6141)to_nested_tensor and from_nested_tensor functionality (#6329, #6330, #6331, #6332)GPSConv Graph Transformer layer and example (#6326, #6327)networkit conversion utilities (#6321)dataset.{attr_name} (#6319)TransE KGE model and example (#6314)FB15k_237 dataset (#3204)Data.update() and HeteroData.update() functionality (#6313)PGExplainer (#6204)AirfRANS dataset (#6287)AttentionExplainer (#6279)LinkNeighborLoader (#6264)BA2MotifDataset explainer dataset (#6257)CycleMotif motif generator to generate n-node cycle shaped motifs (#6256)InfectionDataset to evaluate explanations (#6222)characterization_score and fidelity_curve_auc explainer metrics (#6188)get_message_passing_embeddings (#6201)PointGNNConv layer (#6194)GridGraph graph generator to generate grid graphs (#6220visualize_feature_importance to support node feature visualizations (#6094)Explanation framework (#6091, #6218)CustomMotif motif generator (#6179)ERGraph graph generator to generate Ergos-Renyi (ER) graphs (#6073)BAGraph graph generator to generate Barabasi-Albert graphs - the usage of datasets.BAShapes is now deprecated (#6072seed_time attribute to temporal NodeLoader outputs in case input_time is given (#6196)Data.edge_subgraph and HeteroData.edge_subgraph functionalities (#6193)input_time option to LightningNodeData and transform_sampler_output to NodeLoader and LinkLoader (#6187)summary for PyG/PyTorch models (#5859, #6161)torch.sparse support to PyG (#5906, #5944, #6003, #6033, #6514, #6532, #6748, #6847, #6868, #6874, , , , , , , , , , )inputs_channels back in training benchmark (#6154)utils.to_dense_batch in case max_num_nodes is smaller than the number of nodes (#6124)__dunder__ names (#6999)sort_edge_index, coalesce and to_undirected to only return single edge_index information in case the edge_attr argument is not specified (#6875, #6879, #6893)to_hetero when using an uninitialized submodule without implementing reset_parameters (#6863)get_mesh_laplacian (#6790)GNNExplainer on link prediction tasks (#6787)ChebConv within GNNExplainer (#6778)EdgeStorage.num_edges property (#6710)utils.bipartite_subgraph() and updated docs of HeteroData.subgraph() (#6654)data_list cache of an InMemoryDataset when accessing dataset.data (#6685)Data.subgraph() and HeteroData.subgraph() (#6613)PNAConv and DegreeScalerAggregation to correctly incorporate degree statistics of isolated nodes (#6609)data.to_heterogeneous() filtered attributs in the wrong dimension (#6522)pyg-lib>0.1.0) (#6517)DataLoader workers with affinity to start at cpu0 (#6512)global_*_pool functions (#6504)RGCNConv (#6482)numpy 1.24.0 (#6495)examples/mnist_voxel_grid.py (#6478)LightningNodeData and LightningLinkData code paths (#6473)RGCNConv (#6463)DataParallel class (#6376)ImbalancedSampler on sliced InMemoryDataset (#6374)GraphMultisetTransformer (#6343)transforms.GDC to not crash on graphs with isolated nodes (#6242)InMemoryDataset.data (#6318)SparseTensor dependency in GraphStore (#5517)NeighborSampler with NeighborLoader in the distributed sampling example (#6204)transforms.RemoveIsolatedNodes (#6308)DimeNet that causes a output dimension mismatch (#6305)Data.to_heterogeneous() with empty edge_index (#6304)Explanation.node_mask and Explanation.node_feat_mask (#6267)Explainer to Explanation (#6215)HeteroLinear for un-sorted type vectors (#6198)ExplainerConfig arguments to the Explainer class (#6176)NeighborSampler to be input-type agnostic (#6173)profileit decorator (#6164)GDC example (#6159)torch_geometric.data.lightning (#6140)torch_sparse an optional dependency (#6132, #6134, #6138, #6139)utils.softmax implementation (#6113, #6155, #6805)topk implementation for large enough graphs (#6123)torch-sparse is now an optional dependency (#6625, #6626, #6627, #6628, #6629, #6630)torch-scatter dependencies (#6394, #6395, #6399, #6400, #6615, #6617)GNNExplainer and Explainer from nn.models (#6382)target_index argument in the Explainer interface (#6270)Aggregation.set_validate_args option (#6175)Full Changelog: https://github.com/pyg-team/pytorch_geometric/compare/2.2.0...2.3.0