Neurostorm

NeuroSTORM Skill

Overview

neurostorm-skill wraps the unified NeuroSTORM fMRI platform (CUHK-AIM-Group), which, as of the 2026-05-08 release, ships 8 model implementations under a single training/fine-tuning entry point. Use this skill for the full lifecycle: data download, preprocessing, pretraining, fine-tuning, and inference.

Supported models (8)

Model	Input type	Graph?	Backbone
neurostorm	voxel (4D)	No	Mamba-SSM
swift	voxel (4D)	No	Swin 4D Transformer
braingnn	FC graph (2D)	Yes	GNN
bnt	FC matrix (2D)	No	Transformer
lggnn	ROI + FC	Yes	Learnable GNN
combraintf	FC matrix (2D)	No	Community-aware Transformer
ibgnn	FC graph (2D)	Yes	Interpretable GNN
brainnetcnn	FC matrix (2D)	No	CNN

Supported tasks

ID	Task
1	Age & Gender Prediction
2	Phenotype Prediction
3	Disease Diagnosis
4	fMRI Retrieval
5	Task fMRI State Classification

Supported datasets: HCP1200, ABCD, UKB, Cobre, ADHD200, HCPA, HCPD, UCLA, HCPEP, HCPTASK, GOD, NSD, BOLD5000.

Dual data formats: PT (faster random access, larger disk) and H5 (compact, scales to large cohorts). Choose at preprocessing and at training via --output_format / --data_format.

---

Installation

Use the upstream requirements.txt + set_env.sh flow (Python 3.11, CUDA 12.8, PyTorch 2.7.1).

# 1. Clone and enter
git clone https://github.com/CUHK-AIM-Group/NeuroSTORM.git
cd NeuroSTORM

# 2. Create and activate env
conda create -n neurostorm python=3.11
conda activate neurostorm

# 3. Auto-detect conda + CUDA paths, set TORCH_CUDA_ARCH_LIST
source ./set_env.sh

# 4. Core dependencies
pip install -r requirements.txt
pip install "setuptools<81"               # pytorch-lightning 1.9.4 compat
pip install "transformers<=4.39.3"        # mamba-ssm compat

# 5. Graph-based models (BrainGNN / LG-GNN / IBGNN)
pip install torch-geometric
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.7.0+cu128.html

# 6. FC-based models (BNT / BrainNetCNN / Com-BrainTF)
pip install scikit-learn pandas h5py deepdish

# 7. Mamba-SSM (NeuroSTORM only)
bash scripts/install_mamba.sh
#   or manually: causal-conv1d v1.5.0.post8, mamba v2.2.2
#   both built with TORCH_CUDA_ARCH_LIST matching your GPU (12.0 Blackwell,
#   9.0 H100, 8.9 4090, 8.6 3090, 8.0 A100)

Docker alternative:

docker build -t neurostorm:latest .
docker run --gpus all -it --rm -v $(pwd):/workspace --shm-size=8g neurostorm:latest

Verify:

python -c "import torch; print(torch.cuda.is_available())"
python -c "import torch_geometric; print('PyG OK')"
python -c "from mamba_ssm import Mamba; print('Mamba OK')"
python -c "from models.neurostorm import NeuroSTORM; print('NeuroSTORM OK')"

Full details: upstream INSTALLATION.md.

---

Workflows

1. Data Preprocessing

Assume raw fMRI is in MNI152 space (apply FSL / fMRIPrep / HCP pipelines first).

# 1a. Brain extraction (optional, FSL BET)
bash datasets/brain_extraction.sh /path/to/raw /path/to/extracted

# 1b. Volume preprocessing — 4D voxel tensors for NeuroSTORM / SwiFT
python datasets/preprocessing_volume.py \
    --dataset_name hcp \
    --load_root ./data/hcp \
    --save_root ./processed_data/hcp \
    --output_format pt \                  # or h5 for large cohorts
    --num_processes 8

# 1c. Extract ROI time series — for all graph / FC models
python datasets/generate_roi_data_from_nii.py \
    --atlas_names cc200 \
    --dataset_names hcp \
    --output_dir ./processed_data \
    --num_processes 32

# 1d. Compute functional connectivity — for BrainGNN / BNT / Com-BrainTF / IBGNN / BrainNetCNN
python datasets/compute_fc.py \
    --roi_dir ./processed_data/roi/cc200 \
    --output_dir ./processed_data/fc/cc200 \
    --atlas_name cc200 \
    --fc_types correlation partial_correlation \
    --num_processes 8

Auxiliary scripts: datasets/compute_stats_and_mask.py, datasets/compute_atlas_map.py.

---

2. Pretraining

NeuroSTORM supports two pretraining strategies via main.py.

MAE pretraining (NeuroSTORM):

python main.py \
    --dataset_name HCP1200 \
    --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --model neurostorm \
    --pretraining \
    --use_mae \
    --mask_ratio 0.75 \
    --batch_size 16 \
    --learning_rate 1e-4 \
    --max_epochs 100 \
    --loggername tensorboard \
    --project_name pt_neurostorm_mae

Contrastive pretraining (SwiFT-style):

python main.py \
    --dataset_name HCP1200 \
    --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --model swift \
    --pretraining \
    --use_contrastive \
    --contrastive_type 3 \
    --batch_size 16 \
    --learning_rate 1e-4 \
    --max_epochs 100

Ready-made scripts in scripts/hcp_pretrain/.

---

3. Fine-tuning

The same main.py handles every model; switch with --model and (for graph/FC models) --data_type / --atlas_name / --fc_type / --num_rois.

NeuroSTORM — gender classification:

python main.py \
    --dataset_name HCP1200 \
    --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --model neurostorm \
    --load_model_path ./pretrained_models/neurostorm_mae.pth \
    --downstream_task_type classification \
    --task_name sex \
    --num_classes 2 \
    --batch_size 32 \
    --learning_rate 5e-5 \
    --max_epochs 50

NeuroSTORM — age regression (with label standardization):

python main.py \
    --model neurostorm \
    --downstream_task_type regression \
    --task_name age \
    --num_classes 1 \
    --label_scaling_method standardization \
    --dataset_name HCP1200 --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --batch_size 32 --learning_rate 1e-3 --max_epochs 50

BrainGNN (FC graph input):

python main.py \
    --model braingnn \
    --data_type fc_graph \
    --atlas_name cc200 \
    --fc_type partial_correlation \
    --num_rois 200 \
    --dataset_name HCP1200 --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --downstream_task_type classification --task_name sex --num_classes 2 \
    --batch_size 32

BrainNetworkTransformer (BNT, hierarchical pooling):

python main.py \
    --model bnt \
    --data_type fc_bnt \
    --atlas_name cc200 \
    --num_rois 200 \
    --pooling_sizes 100 50 25 \
    --do_pooling True True False \
    --dataset_name HCP1200 --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --downstream_task_type classification --task_name sex --num_classes 2

BrainNetCNN:

python main.py \
    --model brainnetcnn \
    --data_type fc_bnt \
    --atlas_name cc200 --num_rois 200 \
    --dataset_name HCP1200 --image_path ./data/HCP1200_MNI_to_TRs_minmax \
    --downstream_task_type classification --task_name sex --num_classes 2

LG-GNN, Com-BrainTF, IBGNN: same pattern — set --model and choose the matching --data_type (fc_graph for GNNs, fc_bnt for transformer/CNN FC inputs). See scripts/run_braingnn.sh, scripts/run_bnt.sh, and other scripts/*_downstream/ folders for templates.

Useful fine-tuning flags

Flag	Purpose
--data_format {auto,pt,h5}	select preprocessed file format
--load_model_path	load pretrained backbone weights
--freeze_feature_extractor	freeze backbone, train head only
--resume_ckpt_path	resume from Lightning checkpoint
--use_scheduler --milestones 50 100	multi-step LR
--optimizer AdamW --weight_decay 0.01	switch optimizer
--augment_during_training + --augment_only_affine / --augment_only_intensity	data augmentation
--gpu_ids 0,1,2 / --num_gpus 4	GPU selection (DDP auto when >1)
--loggername tensorboard --project_name NAME	logging

---

4. Inference / Demo

Single subject:

python demo.py \
    --mode single \
    --ckpt_path ./pretrained_models/gender.ckpt \
    --fmri_path ./data/HCP1200_MNI_to_TRs_minmax/img/100206 \
    --task gender

Task options include age, gender, phenotype (with --phenotype_name + --phenotype_type).

Batch inference on a test split:

python demo.py \
    --mode dataset \
    --ckpt_path /path/to/model.ckpt \
    --task age \
    --image_path /path/to/preprocessed/data

Or run the bundled script: sh scripts/run_demo.sh.

---

Input / Output Summary

Stage	Input	Output
Preprocessing (volume)	.nii / .nii.gz in MNI152	.pt or .h5 4D tensors
Preprocessing (ROI)	.nii + atlas	ROI time series .pt/.h5
Preprocessing (FC)	ROI time series	FC matrices (correlation / partial)
Pretraining	Preprocessed voxel tensors	.pth / .ckpt
Fine-tuning	Preprocessed data + pretrained .pth	Fine-tuned .ckpt + TensorBoard logs
Inference	Preprocessed data + .ckpt	Predictions (stdout / file)

---

Testing

Upstream ships a full pytest suite and GitHub Actions CI.

make test           # full suite
make test-cov       # with coverage
make test-unit      # unit tests only
make ci             # local CI dry-run

Key test modules: test_model_loading.py, test_dual_format.py, test_atlas_masking.py.

---

Directory Reference (upstream)

NeuroSTORM/
├── main.py                 entry point for pretraining + fine-tuning
├── demo.py                 unified single-file and dataset inference
├── set_env.sh              auto-detect conda/CUDA paths
├── Makefile                test / dev commands
├── requirements.txt
├── INSTALLATION.md         detailed install
├── USER_GUIDE.md           full usage guide
├── datasets/
│   ├── preprocessing_volume.py
│   ├── generate_roi_data_from_nii.py
│   ├── compute_fc.py
│   ├── fmri_datasets.py    voxel dataset loaders
│   └── roi_datasets.py     ROI + FC loaders
├── models/
│   ├── neurostorm.py  swift.py  braingnn.py  bnt.py
│   ├── lggnn.py  combraintf.py  ibgnn.py  brainnetcnn.py
│   ├── heads/{cls,reg,emb}_head.py
│   ├── load_model.py
│   └── lightning_model.py
├── scripts/
│   ├── hcp_pretrain/  hcp_downstream/
│   ├── install_mamba.sh  run_demo.sh
│   ├── run_braingnn.sh  run_bnt.sh
│   └── dataset_download/
└── tests/                   pytest suite, runs in GitHub Actions CI

---

Reference

• Paper: Towards a General-Purpose Foundation Model for fMRI Analysis, Wang et al., Nature Biomedical Engineering, 2026. https://www.nature.com/articles/s41551-026-01666-y
• Project: https://cuhk-aim-group.github.io/NeuroSTORM/
• GitHub: https://github.com/CUHK-AIM-Group/NeuroSTORM
• Upstream docs: INSTALLATION.md, USER_GUIDE.md

Model attributions: SwiFT (Transconnectome), BrainGNN (LifangHe), BNT (Wayfear), LG-GNN (cnuzh), Com-BrainTF (ubc-tea), IBGNN (HennyJie), BrainNetCNN (nicofarr).

---

Created At: 2026-04-02 00:23 HKT Last Updated At: 2026-05-11 20:45 HKT (synced to upstream 2026-05-08 release: +7 models, dual PT/H5, FC pipeline, pytest suite) Author: chengwang96