This codebase is an extension of Jittor Geometric (https://github.com/AlgRUC/JittorGeometric), specifically developed for the NeutronJittor paper. It is a graph machine learning library based on the Jittor framework, tailored for research and applications in Graph Neural Networks (GNNs).

Note: This is a blind open-source version for the NeutronJittor paper review process.

• Faster GNN Training and Inference: Leveraging Jittor’s dynamic compilation, Jittor Geometric achieves significant speedups in both training and inference.
• Comprehensive Spectral Domain Support: Supports various spectral methods, enabling a wide range of spectral-based GNN architectures.
• Rich Dynamic Dataset Support: Easily handle dynamic datasets, allowing for efficient processing and transformation of graph data.

Jittor Geometric supports a variety of graph datasets, including Cora, CiteSeer, and PubMed. Here’s an example of loading the Cora dataset:

import os.path as osp
from jittor_geometric.datasets import Planetoid
import jittor_geometric.transforms as T

dataset = 'Cora'
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
data = dataset[0]

The following code defines a simple two-layer Graph Convolutional Network (GCN):

from jittor import nn
from jittor_geometric.nn import GCNConv

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_features, 16, cached=True)
        self.conv2 = GCNConv(16, dataset.num_classes, cached=True)

    def execute(self):
        x, edge_index = data.x, data.edge_index
        x = nn.relu(self.conv1(x, edge_index))
        x = nn.dropout(x)
        x = self.conv2(x, edge_index)
        return nn.log_softmax(x, dim=1)

Jittor Geometric includes implementations of popular GNN models, such as:

• Graph Convolutional Network (GCN)
• Chebyshev Network (ChebNet)
• Simplified Graph Convolution (SGC)
• GCNII (Graph Convolutional Network II)

Follow these steps to install Jittor Geometric:

1. Install the Jittor framework by following the Jittor official documentation.
2. Clone this repository and navigate to the project directory:

   git clone <repo_url>
   cd jittor_geometric

3. Install the required dependencies:

   pip install -r requirements.txt

4. When you install the Python package from source code, run the following command:

   git clone https://github.com/AlgRUC/JittorGeometric.git
   cd "the project root directory that contains the setup.py file"
   pip install .

NeutronJittor provides a user-transparent partition training system that allows you to start partition-based GNN training with just one line of code.

• Auto-detection: Automatically detects GPU environment and suggests optimal number of partitions
• User-transparent: Alpha, beta, and sliding window parameters are set to default values (0.5, 0.5, 2) and hidden from users
• Flexible modes: Supports three partition modes: cached, chunk, and optim

python partition_runner.py --dataset cora

python partition_runner.py --dataset cora --num_parts 8

Cached mode (default):

python partition_runner.py --dataset cora --mode cached

Chunk mode:

python partition_runner.py --dataset cora --mode chunk

Optim mode:

python partition_runner.py --dataset cora --mode optim

python partition_runner.py --dataset cora --use_gdc

• cora, citeseer, pubmed
• computers, photo
• chameleon, squirrel
• reddit
• ogbn-arxiv, ogbn-products
• roman_empire, amazon_ratings, minesweeper, questions, tolokers