JittorGeometric 2.0

A Comprehensive Graph Machine Learning Library Built on Jittor

Documentation • Examples • Installation • Quick Start • Models

Overview

JittorGeometric 2.0 is a state-of-the-art graph machine learning library built on the 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.)

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

### Dataset Selection
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 = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]
v_num = data.x.shape[0]

### Data Preprocess
from jittor_geometric.ops import cootocsr,cootocsc
from jittor_geometric.nn.conv.gcn_conv import gcn_norm
edge_index, edge_weight = data.edge_index, data.edge_attr
edge_index, edge_weight = gcn_norm(
                        edge_index, edge_weight,v_num,
                        improved=False, add_self_loops=True)
with jt.no_grad():
   data.csc = cootocsc(edge_index, edge_weight, v_num)
   data.csr = cootocsr(edge_index, edge_weight, v_num)

### Model Definition
from jittor import nn
from jittor_geometric.nn import GCNConv

class GCN(nn.Module):
    def __init__(self, dataset, dropout=0.8):
        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):
        x, csc, csr = data.x, data.csc, data.csr
        x = nn.relu(self.conv1(x, csc, csr))
        x = nn.dropout(x, self.dropout, is_train=self.training)
        x = self.conv2(x, csc, csr)
        return nn.log_softmax(x, dim=1)

### Training
model = GCN(dataset)
optimizer = nn.Adam(params=model.parameters(), lr=0.001, weight_decay=5e-4) 
for epoch in range(200):
   model.train()
   pred = model()[data.train_mask]
   label = data.y[data.train_mask]
   loss = nn.nll_loss(pred, label)
   optimizer.step(loss)

Supported Models

JittorGeometric 2.0 includes implementations of 40+ state-of-the-art GNN models:

Classic Graph Neural Networks

Spectral Graph Neural Networks

Dynamic Graph Neural Networks

Molecular Graph Neural Networks

Graph Self-supervised Learning

Graph Recommendation

Graph Embedding

Graph Transformers

Installation

Step-by-Step Installation

1. Create a conda environment

   conda create -n jittorgeometric python=3.10
   conda activate jittorgeometric
   

2. Install Jittor

   python -m pip install git+https://github.com/Jittor/jittor.git
   

   or follow the Jittor official documentation.

3. Install dependencies

   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 JittorGeometric

   git clone https://github.com/AlgRUC/JittorGeometric.git
   cd JittorGeometric
   pip install .
   

5. Verify installation

   python examples/gcn_example.py
   

For Distributed Training (Optional)

Install MPI support:

conda install -c conda-forge openmpi=4.0.5
conda install -c conda-forge mpi4py

📋 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

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:

python dist_partition.py --dataset reddit --num_parts 2 --use_gdc

3. Launch distributed training:

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.

🧪 Testing

Run a specific example:

python examples/gcn_example.py

📖 Documentation

Comprehensive documentation is available at 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
• Contributors: See Contributors

📄 License

JittorGeometric is released under the Apache 2.0 License.

🙏 Acknowledgments

• Jittor Team for the deep learning framework
• PyTorch Geometric for inspiration
• All contributors and users of JittorGeometric

If you have any questions or would like to contribute, please feel free to contact us!