A lightweight distributed GNN library for full batch node embedding prediction.

Getting started

1. Setup a clean environment.

conda create --name gnn
conda activate gnn

2. Install pytorch (needed for training) and other libraries (needed for downloading datasets).

conda create -n dgl_env python=3.9 -y
conda activate dgl_env

conda install cudatoolkit=11.8 -c nvidia/label/cuda-11.8.0 -y
conda install -c conda-forge nccl

# 安装 PyTorch 2.0+（适配 CUDA 11.8）
conda install pytorch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1  pytorch-cuda=11.8 -c pytorch -c nvidia -y

# 安装 DGL（需指定 CUDA 11.8 版本）
conda install -c dglteam/label/cu118 dgl -y

# 安装cpu支持
conda install pybind11 -y
conda install -c anaconda openmpi -y

3. Compile and install spmm. (Optional. CUDA dev environment needed.)

cd spmm_cpp
python setup.py install

4. Compile and install predictor (Optional. For CPU prediction simulator)

cd predictor_cpp
python setup.py install

5. Prepare datasets (edit the code according to your needs).

//This may take a while.
python prepare_data.py

6. Train.

python main.py

Contact

Contact [email protected] for any problems.