This repository contains the code for SUGAR: Efficient Subgraph-level Training via Resource-aware Graph Partitioning.
Necessary package to start with: torch + dgl + ogb + metis
conda create -n sugar python=3.8
conda activate sugar
# install torch
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
# install dgl
pip install dgl-cu116 dglgo -f https://data.dgl.ai/wheels/repo.html
# (may need this step if 'import dgl' yields error)
conda install cudatoolkit
# install ogb
pip install ogb
# install metis
pip install metis
(1) Download METIS here, and unzip the file.
(2) Build METIS
cd metis-5.1.0
make config shared=1
make install
make
(3) If libmetis.dylib
or libmetis.so
is not under /usr/local/lib
, find it inside directory build
,
and set up the METIS_DLL environment
export METIS_DLL=/usr/local/lib/libmetis.dylib
# or
export METIS_DLL=/your_directory/libmetis.dylib
(4) Check METIS is properly installed
python
>>> import metis
Run baseline GCN with the full graph
python train.py --run-full
Running this command will automatically download the dataset to ./dataset/
.
Run SUGAR:
python train.py --first-time
When running the command for first time, adding '--first-time' will first partition the graph, and then start subgraph-level training.
The graph partitioning results are saved to ./partition
.
Coming soon