From 18742b729faa3b31b66788413b6e65b35c82b667 Mon Sep 17 00:00:00 2001 From: lusz <578752274@qq.com> Date: Sat, 7 Feb 2026 01:48:08 +0800 Subject: [PATCH] gatv2 --- examples/gatv2_example.py | 153 +++++++++++++ jittor_geometric.egg-info/PKG-INFO | 286 ++++++++++++++++++------- jittor_geometric.egg-info/SOURCES.txt | 8 + jittor_geometric/nn/conv/gatv2_conv.py | 92 ++++++++ 4 files changed, 463 insertions(+), 76 deletions(-) create mode 100644 examples/gatv2_example.py create mode 100644 jittor_geometric/nn/conv/gatv2_conv.py diff --git a/examples/gatv2_example.py b/examples/gatv2_example.py new file mode 100644 index 0000000..6d274de --- /dev/null +++ b/examples/gatv2_example.py @@ -0,0 +1,153 @@ +'''Author: lusz +Date: 2024-07-01 15:42:52 +Description: +''' +import os.path as osp +import argparse + +import jittor as jt +from jittor import nn +import sys,os +root = osp.dirname(osp.dirname(osp.abspath(__file__))) +sys.path.append(root) +from jittor_geometric.datasets import Planetoid +import jittor_geometric.transforms as T +from jittor_geometric.nn import GATv2Conv +import time +from jittor import Var +from jittor_geometric.utils import add_remaining_self_loops +from jittor_geometric.utils.num_nodes import maybe_num_nodes +from jittor_geometric.data import CSC,CSR +from jittor_geometric.ops import cootocsr,cootocsc + +# Setup configuration +jt.flags.use_cuda = 1 +jt.flags.lazy_execution = 0 +# jt.misc.set_global_seed(42) +jt.cudnn.set_max_workspace_ratio(0.0) +# Edge normalization for GCN +def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False, + add_self_loops=True, dtype=None): + + fill_value = 2. if improved else 1. + + if isinstance(edge_index, Var): + num_nodes = maybe_num_nodes(edge_index, num_nodes) + + if edge_weight is None: + edge_weight = jt.ones((edge_index.size(1), )) + + if add_self_loops: + edge_index, tmp_edge_weight = add_remaining_self_loops( + edge_index, edge_weight, fill_value, num_nodes) + assert tmp_edge_weight is not None + edge_weight = tmp_edge_weight + + row, col = edge_index[0], edge_index[1] + shape = list(edge_weight.shape) + shape[0] = num_nodes + deg = jt.zeros(shape) + deg = jt.scatter(deg, 0, col, src=edge_weight, reduce='add') + deg_inv_sqrt = deg.pow(-0.5) + deg_inv_sqrt.masked_fill(deg_inv_sqrt == float('inf'), 0) + return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col] + +# Parse arguments +parser = argparse.ArgumentParser() +parser.add_argument('--use_gdc', action='store_true', + help='Use GDC preprocessing.') +parser.add_argument('--dataset', help='graph dataset') +args = parser.parse_args() +dataset=args.dataset + +# Load dataset +path = osp.join(osp.dirname(osp.realpath(__file__)), '../data') +dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures()) +data = dataset[0] + +# Apply GDC preprocessing if requested +if args.use_gdc: + gdc = T.GDC(self_loop_weight=1, normalization_in='sym', + normalization_out='col', + diffusion_kwargs=dict(method='ppr', alpha=0.05), + sparsification_kwargs=dict(method='topk', k=128, + dim=0), exact=True) + data = gdc(data) + +# Prepare data and edge normalization +v_num = data.x.shape[0] +e_num = data.edge_index.shape[1] +edge_index, edge_weight=data.edge_index,data.edge_attr + +edge_index, edge_weight = gcn_norm( + edge_index, edge_weight,v_num, + False, True) + +# Convert to sparse matrix format +with jt.no_grad(): + data.csc = cootocsc(edge_index, edge_weight, v_num) + data.csr = cootocsr(edge_index, edge_weight, v_num) + +# GAT model with two attention layers +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = GATv2Conv(dataset.num_features,128,e_num, cached=True, + normalize=not args.use_gdc) + self.conv2 = GATv2Conv(128, dataset.num_classes,e_num, cached=True, + normalize=not args.use_gdc) + + def execute(self): + x, csc =data.x , data.csc + x = nn.relu(self.conv1(x, csc)) + x = nn.dropout(x) + x = nn.relu(self.conv2(x,csc)) + return nn.log_softmax(x, dim=1) + + +# Initialize model and optimizer +model, data = Net(), data +optimizer = nn.Adam([ + dict(params=model.conv1.parameters(), weight_decay=1e-4), + dict(params=model.conv2.parameters(), weight_decay=1e-4) +], lr=5e-3) + + +# Training function +def train(): + model.train() + + pred = model()[data.train_mask] + label = data.y[data.train_mask] + loss = nn.nll_loss(pred, label) + optimizer.step(loss) + +# Evaluation function +def test(): + model.eval() + logits, accs = model(), [] + # Evaluate on train, val, test sets + for _, mask in data('train_mask', 'val_mask', 'test_mask'): + y_ = data.y[mask] + logits_=logits[mask] + pred, _ = jt.argmax(logits_, dim=1) + acc = pred.equal(y_).sum().item() / mask.sum().item() + accs.append(acc) + return accs + + + +# Training loop +train() +best_val_acc = test_acc = 0 +for epoch in range(1, 201): + train() + train_acc, val_acc, tmp_test_acc = test() + # Track best validation accuracy + if val_acc > best_val_acc: + best_val_acc = val_acc + test_acc = tmp_test_acc + log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}' + print(log.format(epoch, train_acc, best_val_acc, test_acc)) + jt.sync_all() + jt.gc() \ No newline at end of file diff --git a/jittor_geometric.egg-info/PKG-INFO b/jittor_geometric.egg-info/PKG-INFO index 3a3adda..14918c1 100644 --- a/jittor_geometric.egg-info/PKG-INFO +++ b/jittor_geometric.egg-info/PKG-INFO @@ -1,22 +1,51 @@ Metadata-Version: 2.1 Name: jittor_geometric -Version: 0.1 +Version: 2.0.0 +Summary: A comprehensive graph machine learning library built on Jittor +Author: JittorGeometric Team Description-Content-Type: text/markdown License-File: LICENSE +# JittorGeometric 2.0 +

-# JittorGeometric -**[Documentation](https://algruc.github.io/JittorGeometric/index.html)** -JittorGeometric is a graph machine learning library based on the Jittor framework. As a Chinese-developed library, it is tailored for research and applications in Graph Neural Networks (GNNs), aiming to provide an efficient and flexible GNN implementation for researchers and engineers working with graph-structured data. +
+ +[![Documentation](https://img.shields.io/badge/docs-latest-blue.svg)](https://jittor.github.io/JittorGeometric/) +[![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)](LICENSE) +[![Framework](https://img.shields.io/badge/Framework-Jittor-orange.svg)](https://cg.cs.tsinghua.edu.cn/jittor/) +[![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)](https://www.python.org/) + +**A Comprehensive Graph Machine Learning Library Built on Jittor** + +[Documentation](https://algruc.github.io/JittorGeometric/index.html) β€’ [Examples](./examples) β€’ [Installation](#installation) β€’ [Quick Start](#quick-tour) β€’ [Models](#supported-models) + +
+ +--- + +## Overview + +**JittorGeometric 2.0** is a state-of-the-art graph machine learning library built on the [Jittor](https://cg.cs.tsinghua.edu.cn/jittor/) framework. As a Chinese-developed deep learning library, JittorGeometric provides comprehensive support for Graph Neural Networks (GNNs) research and applications, featuring enhanced performance, flexibility, and scalability. + +## 🌟 Key Features + +### Core Capabilities +- **πŸš€ JIT Compilation**: Leverage Just-In-Time compilation for dynamic code modification without pre-compilation overhead +- **⚑ Optimized Sparse Operations**: High-performance sparse matrix computations with CuSparse acceleration +- **🎯 Comprehensive Model Zoo**: 40+ implemented models covering classic, spectral, dynamic, molecular, and transformer-based GNNs +- **πŸ“Š Rich Dataset Support**: Built-in loaders for popular graph datasets (Planetoid, OGB, Reddit, etc.) -## Highlights -- **Easier Code Modification with JIT (Just-In-Time) Compilation**: JittorGeometric leverages JIT compilation to enable easier code modification without any pre-compilation requirements. -- **Optimized Sparse Matrix Computation**: JittorGeometric provides a rich set of operators and utilizes CuSparse to accelerate sparse matrix computations. -- **Comprehensive Domain Support**: Supports both basic and advanced GNNs, including spectral GNNs, dynamic GNNs, and molecular GNNs. +### New in Version 2.0 +- **πŸ”„ Distributed Training**: Multi-GPU and multi-node training support with MPI +- **🌊 Dynamic Graph Processing**: Event-based dynamic graph support with parallel processing +- **πŸ“¦ Mini-batch Support**: Efficient mini-batch training for large-scale graphs +- **πŸ”§ Ascend-GNN**: GNN for NPU +- **πŸŽ›οΈ Extended Model Categories**: Graph transformers, self-supervised learning, and recommendation systems ## Quick Tour @@ -25,11 +54,13 @@ JittorGeometric is a graph machine learning library based on the Jittor framewor import os.path as osp from jittor_geometric.datasets import Planetoid import jittor_geometric.transforms as T +import jittor as jt -dataset = 'Cora' -path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset) +dataset = 'cora' +path = osp.join(osp.dirname(osp.realpath(__file__)), '.', 'data', dataset) dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures()) data = dataset[0] +v_num = data.x.shape[0] ### Data Preprocess from jittor_geometric.ops import cootocsr,cootocsc @@ -48,9 +79,9 @@ from jittor_geometric.nn import GCNConv class GCN(nn.Module): def __init__(self, dataset, dropout=0.8): - super(Net, self).__init__() - self.conv1 = GCNConv(in_channels=dataset.num_features, out_channels=256,spmm=args.spmm) - self.conv2 = GCNConv(in_channels=256, out_channels=dataset.num_classes,spmm=args.spmm) + super(GCN, self).__init__() + self.conv1 = GCNConv(in_channels=dataset.num_features, out_channels=256) + self.conv2 = GCNConv(in_channels=256, out_channels=dataset.num_classes) self.dropout = dropout def execute(self): @@ -73,102 +104,205 @@ for epoch in range(200): ``` ## Supported Models -JittorGeometric includes implementations of popular GNN models, such as: -### Classic Graph Neural Networks +JittorGeometric 2.0 includes implementations of 40+ state-of-the-art GNN models: -| Model | Year | Venue | -|------------|------|--------| -| [ChebNet](./examples/chebnet_example.py) | 2016 | NeurIPS| -| [GCN](./examples/gcn_example.py) | 2017 | ICLR | -| [GraphSAGE](./examples/graphsage_example.py) | 2017 | NeurIPS| -| [GAT](./examples/gat_example.py) | 2018 | ICLR | -| [SGC](./examples/sgc_example.py) | 2019 | ICML | -| [APPNP](./examples/appnp_example.py) | 2019 | ICLR | -| [GCNII](./examples/gcn2_example.py) | 2020 | ICML | +### Classic Graph Neural Networks ---- +| Model | Year | Venue | Description | +|------------|------|--------|-------------| +| [ChebNet](./examples/chebnet_example.py) | 2016 | NeurIPS| Spectral graph convolutions | +| [GCN](./examples/gcn_example.py) | 2017 | ICLR | Graph Convolutional Networks | +| [GraphSAGE](./examples/graphsage_example.py) | 2017 | NeurIPS| Inductive graph learning | +| [GAT](./examples/gat_example.py) | 2018 | ICLR | Graph Attention Networks | +| [SGC](./examples/sgc_example.py) | 2019 | ICML | Simplified Graph Convolution | +| [APPNP](./examples/appnp_example.py) | 2019 | ICLR | Approximate Personalized Propagation | +| [GCNII](./examples/gcn2_example.py) | 2020 | ICML | Deeper Graph Convolutional Networks | ### Spectral Graph Neural Networks -| Model | Year | Venue | -|------------|------|--------| -| [GPRGNN](./examples/gprgnn_example.py) | 2021 | ICLR | -| [BernNet](./examples/bernnet_example.py) | 2021 | NeurIPS| -| [ChebNetII](./examples/chebnet2_example.py) | 2022 | NeurIPS| -| [EvenNet](./examples/evennet_example.py) | 2022 | NeurIPS| -| [OptBasis](./examples/optbasis_example.py) | 2023 | ICML | - ---- +| Model | Year | Venue | Description | +|------------|------|--------|-------------| +| [GPRGNN](./examples/gprgnn_example.py) | 2021 | ICLR | Generalized PageRank GNN | +| [BernNet](./examples/bernnet_example.py) | 2021 | NeurIPS| Bernstein polynomial filters | +| [ChebNetII](./examples/chebnet2_example.py) | 2022 | NeurIPS| Improved Chebyshev filters | +| [EvenNet](./examples/evennet_example.py) | 2022 | NeurIPS| Even polynomial filters | +| [OptBasis](./examples/optbasis_example.py) | 2023 | ICML | Optimal basis functions | ### Dynamic Graph Neural Networks -| Model | Year | Venue | -|------------|------|--------| -| [JODIE](./examples/jodie_example.py) | 2019 | SIGKDD | -| [DyRep](./examples/dyrep_example.py) | 2019 | ICLR | -| [TGN](./examples/tgn_example.py) | 2020 | ArXiv | -| [GraphMixer](./examples/graphmixer_example.py) | 2022 | ICLR | -| [Dygformer](./examples/dygformer_example.py) | 2023 | NeurIPS| - ---- +| Model | Year | Venue | Description | +|------------|------|--------|-------------| +| [JODIE](./examples/jodie_example.py) | 2019 | SIGKDD | Temporal interaction networks | +| [DyRep](./examples/dyrep_example.py) | 2019 | ICLR | Dynamic representation learning | +| [TGN](./examples/tgn_example.py) | 2020 | ArXiv | Temporal Graph Networks | +| [GraphMixer](./examples/graphmixer_example.py) | 2022 | ICLR | MLP-based dynamic graphs | +| [Dygformer](./examples/dygformer_example.py) | 2023 | NeurIPS| Dynamic graph transformers | ### Molecular Graph Neural Networks -| Model | Year | Venue | -|------------|------|--------| -| [SchNet](./examples/schnet_example.py) | 2017 | NeurIPS| -| [DimeNet](./examples/dimenet_example.py) | 2020 | ICLR | -| [EGNN](./examples/egnn_example.py) | 2021 | ICML | -| [SphereNet](./examples/spherenet_example.py) | 2022 | ICLR | -| [Uni-Mol](./examples/unimol_example.py) | 2023 | ICLR | - ---- +| Model | Year | Venue | Description | +|------------|------|--------|-------------| +| [SchNet](./examples/schnet_example.py) | 2017 | NeurIPS | Continuous-filter convolutions | +| [DimeNet](./examples/dimenet_example.py) | 2020 | ICLR | Directional message passing | +| [EGNN](./examples/egnn_example.py) | 2021 | ICML | Equivariant Graph Networks | +| [Graphormer](./examples/graphormer_example.py) | 2021 | NeurIPS | Graph transformers for molecules | +| [SphereNet](./examples/spherenet_example.py) | 2022 | ICLR | Spherical message passing | +| [Uni-Mol](./examples/unimol_example.py) | 2023 | ICLR | Universal molecular representation | +| [Transformer-M](./examples/transformer-m_example.py) | 2023 | ICLR | Molecular transformers | ### Graph Self-supervised Learning -| Model | Year | Venue | -|------------------------------------------|------|-------| -| [DGI](./examples/dgi_example.py) | 2019 | ICLR | -| [MVGRL](./examples/mvgrl_example.py) | 2020 | ICML | -| [GRACE](./examples/grace_example.py) | 2020 | ICML | -| [PolyGCL](./examples/polygcl_example.py) | 2024 | ICLR | +| Model | Year | Venue | Description | +|------------------------------------------|------|-------|-------------| +| [DGI](./examples/dgi_example.py) | 2019 | ICLR | Deep Graph Infomax | +| [MVGRL](./examples/mvgrl_example.py) | 2020 | ICML | Multi-view contrastive learning | +| [GRACE](./examples/grace_example.py) | 2020 | ICML | Graph contrastive learning | +| [PolyGCL](./examples/polygcl_example.py) | 2024 | ICLR | Polynomial graph contrastive learning | + +### Graph Recommendation + +| Model | Year | Venue | Description | +|------------------------------------------|------|-------|-------------| +| [SASREC](./examples/sasrec_example.py) | 2018 | ICDM | Self-attentive sequential recommendation | +| [SGNNHN](./examples/sgnnhn_example.py) | 2020 | CIKM | Set-based GNN for heterogeneous networks | +| [CRAFT](./examples/craft_example.py) | 2025 | ArXiv | Cross-attention recommendation | + +### Graph Embedding + +| Model | Year | Venue | Description | +|------------------------------------------|------|-------|-------------| +| [Deepwalk](./examples/link_pred_example.py) | 2014 | KDD | Random walk embeddings | +| [LINE](./examples/link_pred_example.py) | 2015 | WWW | Large-scale information network embedding | +| [Node2Vec](./examples/link_pred_example.py) | 2016 | KDD | Scalable feature learning | +| [LightGCN](./examples/recsys_example.py) | 2020 | SIGIR | Simplified GCN for recommendation | +| [DirectAU](./examples/recsys_example.py) | 2022 | KDD | Direct alignment and uniformity | +| [SimGCL](./examples/recsys_example.py) | 2024 | KAIS | Simple graph contrastive learning | +| [XSimGCL](./examples/recsys_example.py) | 2024 | TKDE | Extreme simple graph contrastive learning | + +### Graph Transformers + +| Model | Year | Venue | Description | +|------------------------------------------|------|-------|-------------| +| [SGFormer](./examples/sgformer_example.py) | 2023 | NeurIPS | Simplifying graph transformers | +| [NAGFormer](./examples/nagphormer_example.py) | 2023 | ICLR | Neighborhood aggregation transformers | +| [PolyFormer](./examples/polyformer_example.py) | 2024 | KDD | Polynomial-based graph transformers | ## Installation -Follow these steps to install JittorGeometric: -1. Create a conda environment +### Step-by-Step Installation + +1. **Create a conda environment** ```bash - conda create -n jittorGeometric python=3.10 - conda activate jittorGeometric + conda create -n jittorgeometric python=3.10 + conda activate jittorgeometric ``` -2. Install Jittor: +2. **Install Jittor** ```bash python -m pip install git+https://github.com/Jittor/jittor.git ``` - or by following the [Jittor official documentation](https://cg.cs.tsinghua.edu.cn/jittor/). -3. Installing other dependencies: - ```bash - pip install astunparse==1.6.3 autograd==1.7.0 cupy==13.3.0 numpy==1.24.0 pandas==2.2.3 Pillow==11.1.0 PyMetis==2023.1.1 six==1.16.0 pyparsing==3.2 scipy==1.15.1 setuptools==69.5.1 sympy==1.13.3 tqdm==4.66.4 einops huggingface_hub==0.27.1 + or follow the [Jittor official documentation](https://cg.cs.tsinghua.edu.cn/jittor/). +3. **Install dependencies** + ```bash + pip install astunparse==1.6.3 autograd==1.7.0 cupy==13.3.0 numpy==1.24.0 \ + pandas==2.2.3 Pillow==11.1.0 PyMetis==2023.1.1 six==1.16.0 \ + pyparsing==3.2 scipy==1.15.1 setuptools==69.5.1 sympy==1.13.3 \ + tqdm==4.66.4 einops huggingface_hub==0.27.1 networkx==3.4.2 \ + scikit-learn==1.7.1 rdkit==2025.3.5 seaborn==0.13.2 \ + alive-progress==3.3.0 ``` -4. Install the package: + +4. **Install JittorGeometric** ```bash git clone https://github.com/AlgRUC/JittorGeometric.git cd JittorGeometric pip install . ``` -5. Verify the installation by running the gcn_example.py +5. **Verify installation** + ```bash + python examples/gcn_example.py + ``` -## Warnings -Since JittorGeometric is still under development, please note the following: -1. rdkit is temporarily not supported and will be provided in future updates. -2. Users need to configure the cupy environment. +### For Distributed Training (Optional) -## Contributors +Install MPI support: +```bash +conda install -c conda-forge openmpi=4.0.5 +conda install -c conda-forge mpi4py +``` -This project is actively maintained by the **JittorGeometric Team**. +## πŸ“‹ Requirements + +- Python 3.10+ +- CUDA 11.0+ (for GPU support) +- Jittor 1.3.0+ +- CuPy (for CUDA operations) +- NumPy, SciPy, NetworkX +- For distributed training: OpenMPI 4.0.5+ and mpi4py + +## πŸ”„ Distributed Training + +JittorGeometric 2.0 supports distributed training across multiple GPUs and nodes: + +### Single Machine Multi-GPU +```bash +mpiexec -n 2 python dist_gcn.py --num_parts 2 --dataset reddit +``` + +### Multi-Node Training +1. Configure your hostfile: +``` +172.31.195.15 slots=1 +172.31.195.16 slots=1 +``` + +2. Partition the graph: +```bash +python dist_partition.py --dataset reddit --num_parts 2 --use_gdc +``` + +3. Launch distributed training: +```bash +mpirun -n 2 --hostfile hostfile \ +--prefix /path/to/conda/env \ +python dist_gcn.py --num_parts 2 --dataset reddit +``` + +For detailed distributed training setup, see [examples/README.md](./examples/README.md). + +## πŸ§ͺ Testing +Run a specific example: +```bash +python examples/gcn_example.py +``` + +## πŸ“– Documentation + +Comprehensive documentation is available at [https://algruc.github.io/JittorGeometric/index.html](https://algruc.github.io/JittorGeometric/index.html) + +## πŸ‘₯ Contributors + +This project is actively maintained by the **JittorGeometric Team** at Renmin University of China and Northeastern Universityβ€Œ. + +### Core Team +- Project Lead: [runlin_lei@ruc.edu.cn](mailto:runlin_lei@ruc.edu.cn) +- Contributors: See [Contributors](https://github.com/AlgRUC/JittorGeometric/graphs/contributors) + +## πŸ“„ License + +JittorGeometric is released under the [Apache 2.0 License](LICENSE). + +## πŸ™ Acknowledgments + +- [Jittor Team](https://cg.cs.tsinghua.edu.cn/jittor/) for the deep learning framework +- [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/) for inspiration +- All contributors and users of JittorGeometric + +--- -If you have any questions or would like to contribute, please feel free to contact [runlin_lei@ruc.edu.cn](mailto:runlin_lei@ruc.edu.cn). +
+ If you have any questions or would like to contribute, please feel free to contact us! +
diff --git a/jittor_geometric.egg-info/SOURCES.txt b/jittor_geometric.egg-info/SOURCES.txt index 5b628fb..af43039 100644 --- a/jittor_geometric.egg-info/SOURCES.txt +++ b/jittor_geometric.egg-info/SOURCES.txt @@ -13,6 +13,7 @@ jittor_geometric/data/conformer.py jittor_geometric/data/data.py jittor_geometric/data/dataset.py jittor_geometric/data/dictionary.py +jittor_geometric/data/distchunk.py jittor_geometric/data/download.py jittor_geometric/data/graphchunk.py jittor_geometric/data/in_memory_dataset.py @@ -75,6 +76,7 @@ jittor_geometric/nn/conv/even_conv.py jittor_geometric/nn/conv/gat_conv.py jittor_geometric/nn/conv/gcn2_conv.py jittor_geometric/nn/conv/gcn_conv.py +jittor_geometric/nn/conv/gcnacl_conv.py jittor_geometric/nn/conv/gpr_conv.py jittor_geometric/nn/conv/lightgcn_conv.py jittor_geometric/nn/conv/message_passing.py @@ -132,6 +134,7 @@ jittor_geometric/ops/aggregateWithWeight.py jittor_geometric/ops/cootocsc.py jittor_geometric/ops/cootocsr.py jittor_geometric/ops/csctocsr.py +jittor_geometric/ops/distspmm.py jittor_geometric/ops/edgesoftmax.py jittor_geometric/ops/getweight.py jittor_geometric/ops/gpuinitco.py @@ -151,6 +154,8 @@ jittor_geometric/ops/cpp/cootocsr_op.cc jittor_geometric/ops/cpp/cootocsr_op.h jittor_geometric/ops/cpp/csctocsr_op.cc jittor_geometric/ops/cpp/csctocsr_op.h +jittor_geometric/ops/cpp/distspmm_op.cc +jittor_geometric/ops/cpp/distspmm_op.h jittor_geometric/ops/cpp/edgesoftmax_op.cc jittor_geometric/ops/cpp/edgesoftmax_op.h jittor_geometric/ops/cpp/edgesoftmaxbackward_op.cc @@ -173,6 +178,8 @@ jittor_geometric/ops/cpp/toundirected_op.cc jittor_geometric/ops/cpp/toundirected_op.h jittor_geometric/partition/__init__.py jittor_geometric/partition/chunk_manager.py +jittor_geometric/partition/dist_partition.py +jittor_geometric/partition/distchunk_manager.py jittor_geometric/partition/partition_graph.py jittor_geometric/transforms/__init__.py jittor_geometric/transforms/normalize_features.py @@ -181,6 +188,7 @@ jittor_geometric/utils/bprloss.py jittor_geometric/utils/coalesce.py jittor_geometric/utils/convert.py jittor_geometric/utils/degree.py +jittor_geometric/utils/distributed_utils.py jittor_geometric/utils/get_laplacian.py jittor_geometric/utils/gssl_utils.py jittor_geometric/utils/induced_graph.py diff --git a/jittor_geometric/nn/conv/gatv2_conv.py b/jittor_geometric/nn/conv/gatv2_conv.py new file mode 100644 index 0000000..eaf718f --- /dev/null +++ b/jittor_geometric/nn/conv/gatv2_conv.py @@ -0,0 +1,92 @@ +'''Description: +Author: lusz +Date: 2024-06-26 10:57:06 +''' +from typing import Optional, Tuple +from jittor_geometric.typing import Adj, OptVar + +import jittor as jt +from jittor import Var +from jittor_geometric.nn.conv import MessagePassingNts +from jittor_geometric.utils import add_remaining_self_loops +from jittor_geometric.utils.num_nodes import maybe_num_nodes + +from ..inits import glorot, zeros +from jittor_geometric.data import CSC,CSR +from jittor_geometric.ops import ScatterToEdge,EdgeSoftmax,aggregateWithWeight,ScatterToVertex + + +class GATv2Conv(MessagePassingNts): + r"""The graph convolutional operator from the "`Graph Attention Networks V2`" paper. + """ + + _cached_edge_index: Optional[Tuple[Var, Var]] + _cached_csc: Optional[CSC] + + def __init__(self, in_channels: int, out_channels: int,e_num: int, + improved: bool = False, cached: bool = False, + add_self_loops: bool = True, normalize: bool = True, + bias: bool = True, **kwargs): + + kwargs.setdefault('aggr', 'add') + super(GATv2Conv, self).__init__(**kwargs) + + self.in_channels = in_channels + self.out_channels = out_channels + self.improved = improved + self.cached = cached + self.add_self_loops = add_self_loops + self.normalize = normalize + + self._cached_edge_index = None + self._cached_adj_t = None + + self.weight = jt.random((in_channels, out_channels)) + self.edge_weight=jt.random((2*out_channels,1)) + self.reset_parameters() + + def reset_parameters(self): + glorot(self.weight) + glorot(self.edge_weight) + self._cached_adj_t = None + self._cached_csc=None + + def execute(self, x: Var, csc: CSC) -> Var: + """""" + out=self.vertex_forward(x) + out = self.propagate(x=out,csc=csc) + return out + + def propagate(self,x,csc): + e_msg=self.scatter_to_edge(x,csc) + out = self.edge_forward(e_msg,csc) + out=self.scatter_to_vertex(out,csc) + return out + + def scatter_to_edge(self,x,csc)->Var: + out1=ScatterToEdge(x,csc,"src") + out2=ScatterToEdge(x,csc,"dst") + out =jt.contrib.concat([out1,out2],dim=1) + return out + + def edge_forward(self,x,csc)->Var: + m=jt.nn.leaky_relu(x,scale=0.2) + out = m @ self.edge_weight + a=EdgeSoftmax(out,csc) + half_dim=int(jt.size(x,1)/2) + e_msg=x[:,0:half_dim] + return e_msg * a + + def scatter_to_vertex(self,edge,csc)->Var: + out=ScatterToVertex(edge,csc,'src') + return out + + def vertex_forward(self,x:Var)->Var: + out = x @ self.weight + out=jt.nn.relu(out) + return out + + + def __repr__(self): + return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, + self.out_channels) \ No newline at end of file