Skip to content

biomedicalinformaticsgroup/gnn-suite

BranchesTags
 
 

Repository files navigation

GNN-suite

GNN-Suite is a robust and modular framework for constructing and benchmarking Graph Neural Network (GNN) architectures for computational biology applications. Built on the Nextflow scientific workflow system, it standardises experimentation and reproducibility for evaluating GNN model performance.

The framework supports binary classification, multi-class classification, and regression tasks, and integrates MLflow for experiment tracking, metric logging, and model versioning. Automated hyperparameter optimisation is provided via Optuna with YAML-based configuration of search spaces, samplers, and pruners.

To illustrate its utility, we applied the system to the identification of cancer-driver genes by constructing molecular networks from protein-protein interaction (PPI) data sourced from STRING and BioGRID and annotating nodes with features extracted from PCAWG, PID, and COSMIC-CGC repositories. Our experiments showed that all GNN architectures consistently outperformed a baseline logistic regression model, with GCN2 achieving the highest balanced accuracy (0.807 ± 0.035) on a STRING-based network.

Paper

If you use GNN-Suite in your research, please cite the following preprint:

GNN-Suite: a Graph Neural Network Benchmarking Framework for Biomedical Informatics
Sebestyén Kamp, Giovanni Stracquadanio, T. Ian Simpson
arXiv:2505.10711, 2025
DOI: 10.48550/arXiv.2505.10711

@article{kamp2025gnnsuite,
  title     = {GNN-Suite: a Graph Neural Network Benchmarking Framework for Biomedical Informatics},
  author    = {Kamp, Sebestyén and Stracquadanio, Giovanni and Simpson, T. Ian},
  journal   = {arXiv preprint arXiv:2505.10711},
  year      = {2025},
  url       = {https://arxiv.org/abs/2505.10711},
  doi       = {10.48550/arXiv.2505.10711}
}

Models

The following models are included:

  • Graph Convolutional Networks (GCN)
  • Graph Attention Networks (GAT)
  • Hierarchical Graph Convolutional Networks (HGCN)
  • Parallel Hierarchical Graph Convolutional Networks (PHGCN)
  • Graph SAmpling and aggreGatE (GraphSAGE)
  • Graph Transformer Networks (GTN)
  • Graph Isomorphism Networks (GIN)
  • Graph Convolutional Networks II (GCNII)

Task Types

The pipeline supports three types of node classification/prediction tasks:

  • Binary Classification (binary): Two class classification
  • Multiclass Classification (multiclass): Multi class classification with more than two classes
  • Regression (regression): Continuous value prediction

Specify the task type using the --task_type parameter:

# Binary classification
nextflow run main.nf -profile docker,test --task_type binary

# Multiclass classification
nextflow run main.nf -profile docker \
  --geneFile testdata/sfari_multiclass_genes.csv \
  --networkFile testdata/sfari_multiclass_network.tsv \
  --task_type multiclass

# Regression
nextflow run main.nf -profile docker \
  --geneFile testdata/sfari_regression_genes.csv \
  --networkFile testdata/sfari_regression_network.tsv \
  --task_type regression

Architecture

New Architecture

Running the workflow

Install or update the workflow

nextflow pull biomedicalinformaticsgroup/gnn-suite

Run a test

nextflow run biomedicalinformaticsgroup/gnn-suite -profile docker,test

Run an experiment

nextflow run biomedicalinformaticsgroup/gnn-suite -profile docker,<experiment_file>

The results of the experimetn will be stored in the results/data/<experiment_file>/ and results/figures/<experiment_file>/ directory.

For more information on Nextflow, you can visit the official documentation at nextflow.io/docs.

Parameters Reference

Data Parameters

Parameter Value Description
--geneFile null Path to gene features CSV file
--networkFile null Path to network edges TSV file
--resultsDir null Results directory (defaults to ./results)
--dataSet null Dataset name (auto generated from geneFile if not provided genefile_tasktype)

Experimental Parameters

Parameter Value Description
--models ["gcn2", "gcn", "gat"] List of GNN models to train
--epochs [100] Number of training epochs
--learning_rate 0.01 Learning rate
--weight_decay 1e-4 Weight decay for regularization
--train_size 0.8 Train/test split ratio
--replicates 1 Number of experiment replicates
--verbose_interval 10 Logging interval (epochs)
--dropout 0.2 Dropout rate
--alpha 0.1 Alpha parameter (GCNII)
--theta 1 Theta parameter (GCNII)
--num_heads 1 Number of attention heads (GAT)
--task_type binary Task type: binary, multiclass, regression

Evaluation Parameters

Parameter Value Description
--enable_plots false Generate PDF plots

Output Description

The --dataSet parameter determines the output folder name. If not provided, it defaults to <geneFile_basename>_<task_type> (e.g., genes_binary, genes_multiclass).

After running the pipeline, results are organized in the following structure:

results/
├── data/<dataSet>/
│   ├── full-<model>-<epochs>-run-<replicate>-<task_type>.txt        # Full training output
│   ├── full-<model>-<epochs>-run-<replicate>-<task_type>-train.txt  # Training set metrics
│   ├── full-<model>-<epochs>-run-<replicate>-<task_type>-test.txt   # Test set metrics
│   └── full-<model>-<epochs>-run-<replicate>-<task_type>-all.txt    # All nodes metrics
├── hyperparameters/<dataSet>/
│   └── best_trial_<model>_<dataSet>.json    # Optimized hyperparameters from hyperopt
└── figures/<dataSet>/
    └── <model>-<epochs>-split-<split>-<task_type>.pdf   # Training plots (if enabled)

Docker Image

The gnn-suite Docker image is available on GitHub Container Registry.

Pulling the Image

docker pull ghcr.io/essharom/gnn-suite:latest
# or specific version
docker pull ghcr.io/essharom/gnn-suite:v1.0.0

Image Contents

The Docker image (8.6GB) includes:

  • PyTorch 2.2.1 with CUDA 12.1 support
  • PyTorch Geometric 2.7.0 with all extension packages (torch-scatter, torch-sparse, torch-cluster, torch-spline-conv)
  • MLflow 3.11.1 for experiment tracking
  • Optuna for hyperparameter optimization
  • Scientific Python stack: numpy, pandas, scikit-learn, matplotlib, seaborn
  • Development tools: Jupyter, black, pylint, git

Using the Container Directly

Run Python scripts with the container:

docker run --rm -v $(pwd):/workspace -w /workspace \
  ghcr.io/essharom/gnn-suite:latest \
  python bin/gnn.py --help

Viewing MLflow Results

Start MLflow UI to view experiment results from your pipeline runs:

# Start MLflow UI server using the container
docker run -d --rm \
  -p 5000:5000 \
  -v $(pwd)/work/<task_dir>/mlruns:/mlruns \
  ghcr.io/essharom/gnn-suite:latest \
  mlflow ui --backend-store-uri file:///mlruns --host 0.0.0.0 --port 5000

# Access UI at http://localhost:5000

For file-based tracking (default), the MLflow runs are stored in the Nextflow work directory. To view all experiments:

# View experiments from results directory
docker run -d --rm \
  -p 5000:5000 \
  -v $(pwd)/mlruns:/mlruns \
  ghcr.io/essharom/gnn-suite:latest \
  mlflow ui --backend-store-uri file:///mlruns --host 0.0.0.0 --port 5000

Adding a New Experiment

  1. Create a Config File: Create a new configuration file <experiment_file>.config with the parameters for the experiment:

    // profile to test the string workflow
params {
  resultsDir = "${baseDir}/results/"
  networkFile = "${baseDir}/data/<network_file>.tsv"
  geneFile = "${baseDir}/data/<feature_file>.csv"
  epochs = [300]
  models = ["gcn2", "gcn", "gat", "gat3h", "hgcn", "phgcn", "sage", "gin", "gtn"]
  replicates = 10
  verbose_interval = 1
  dropout = 0.2
  alpha = 0.1
  theta = 1
  dataSet = "<experiment_file_tag>"
}

  2. Update base.config: Add a new profile for your experiment in base.config:

    profiles {
  // existing profiles...

  // test profile for the biogrid cosmic network defining some data
  <config_file> {
    includeConfig '<experiment_file>.config'
  }
}

  3. Run the Experiment: Execute the pipeline with the new profile using:

    nextflow run main.nf -profile docker,<experiment_file>

    or

    nextflow run biomedicalinformaticsgroup/gnn-suite -profile docker,<experiment_file>

Adding a New Model

  1. Create Model: Implement the new model class in models.py:

    class NewModel(torch.nn.Module):
    def __init__(self, num_features, num_classes, num_hidden=16, num_layers=2, dropout=0.5):
        super(NewModel, self).__init__()
        # Define layers
    def forward(self, data):
        # Define forward pass

  2. Import Model: Add your model to the imports in gnn.py:

    from models import GCN, GAT, ..., NewModel

  3. Update build_model: Add your model to the build_model function in gnn.py:

    elif name == "new_model":
    return NewModel(num_features, num_classes, dropout=dropout)

  4. Include in Experiment: Add the new model name to the models list in your experiment config (<experiment_file>.config):

    models = ["gcn", "gat", ..., "new_model"]

Hyperparameter Optimization with Optuna

To run the hyperparameter optimization workflow using optuna defined in hyperopt.py:

# Binary classification
nextflow run main.nf -profile docker -entry hyperopt \
  --geneFile testdata/genes.csv \
  --networkFile testdata/network.tsv \
  --task_type binary

# Multiclass classification
nextflow run main.nf -profile docker -entry hyperopt \
  --geneFile testdata/sfari_multiclass_genes.csv \
  --networkFile testdata/sfari_multiclass_network.tsv \
  --task_type multiclass

# Regression
nextflow run main.nf -profile docker -entry hyperopt \
  --geneFile testdata/sfari_regression_genes.csv \
  --networkFile testdata/sfari_regression_network.tsv \
  --task_type regression

Search Space Configuration

The hyperparameter search space is defined in conf/hyperparams.yaml:

Parameter Range Type
learning_rate 0.001 - 0.5 float (log scale)
weight_decay 0.00001 - 0.5 float (log scale)
dropout 0.0 - 0.8 float
epochs 100 - 300 int

Model-Specific Parameters

Model Parameter Range
GCNII alpha 0.001 - 10.0 (log scale)
GCNII theta 0.001 - 10.0 (log scale)
GAT num_heads [1, 2, 4, 8]

Optimization Settings

Setting Value
n_trials 5
sampler TPE

Output

The results are stored in results/hyperparameters/<dataSet>/:

  • best_trial_<model>_<dataSet>.json - Best hyperparameters (auto used during training)

When running the main training workflow, optimized hyperparameters are automatically used if available.

MLflow Integration

The pipeline supports MLflow for experiment tracking, metrics logging, and model registry.

Enabling MLflow

# Using SQLite database
nextflow run main.nf -profile docker,test \
  --with_mlflow true \
  --mlflow_tracking_uri "sqlite:///mlflow.db"

MLflow Parameters

Parameter Default Description
--with_mlflow false Enable MLflow tracking
--mlflow_tracking_uri file:./mlruns MLflow tracking URI (sqlite:///mlflow.db or http://localhost:5000)
--mlflow_experiment_name gnn-suite-default Experiment name in MLflow
--mlflow_register_model false Register trained models in MLflow Model Registry

Using a Tracking Server

  1. Start the MLflow server using the provided script:

    bash mlflow_server.sh

    This starts a server with SQLite backend (sqlite:///mlflow.db) and artifact storage (./mlflow-artifacts).

  2. Run the pipeline with the server URI:

    nextflow run main.nf -profile docker,test \
  --with_mlflow true \
  --mlflow_tracking_uri "http://localhost:5000"

  3. View experiments at http://localhost:5000

Logged Metrics

MLflow logs the following metrics per epoch:

  • Binary/Multiclass: precision, recall, accuracy, balanced accuracy, F1, AUC
  • Regression: MSE, RMSE, MAE, R²

FAQ

If you encounter the following error message when attempting to execute the script:

Command error:
  .command.sh: line 2: ../gnn-suite/bin/plot.py: Permission denied

You need to grant the necessary execution permissions to the specific python scripts. You can do this by running (e.g. plot.py):

 chmod +x /home/<path>/code/gnn-suite/bin/plot.py

Authors

  • Sebestyén Kamp
  • Ian Simpson
  • Giovanni Stracquadanio

About

Nextflow pipeline to train several GNN architectures

Resources

Readme
Activity
Custom properties

Stars

3 stars

Watchers

0 watching

Forks

1 fork
Report repository

Releases

1 tags

Packages

 
 
 

Contributors

Languages

  • Jupyter Notebook 60.2%
  • Python 34.7%
  • Nextflow 4.0%
  • Other 1.1%