You need to agree to share your Hugging Face information to access this model
This repository is publicly accessible, but you have to accept the conditions to access its files and content.
By requesting access you agree to the license terms and will use this model for non-commercial research only.
Log in or Sign Up to review the conditions and access this model content.
Brain2Vec: A Self-Supervised 3D Foundation Model for Structural Brain MRI
This is the hugging face model card for Brain2Vec, a 3D foundation model pre-trained on 74,425 T1-weighted brain MRI scans from the iSTAGING consortium.
If you are looking for the preprint, it is located here: Brain2Vec-preprint.
If you are looking for the github repo, it is located here: Brain2Vec-repo.
Gaining access
We suggest to either apply for access on huggingface and use our model directly (i.e. for fine-tuning, ablation studies, etc.) or use our model on NiChart for simpler tasks. We want our model to help the neuroimaging community (researchers, students, etc.), in order to gain access, please mention your university/company, PI and the purpose of usage.
Once you gain access, you can use our model using huggingface hub:
from Brain2Vec.models.brain2vec import brain2vec_base_patch8_dec512d8b
from Brain2Vec.models.vit import vit_base_patch8
from Brain2Vec.models.utils import load_brain2vec
# for Brain2Vec
model = load_brain2vec(
repo_id="cbica/Brain2Vec",
weights_filename="brain2vec.safetensors",
model_class=brain2vec_base_patch8_dec512d8b,
model_kwargs={"img_size": 128},
device="cuda",
)
model.eval()
# you can review inference.py on github to see how to perform inference correctly. We
# use the CLS token for feature extraction, unlike other foundation models that use pooling.
# for ViT-BA
model = load_brain2vec(
repo_id="cbica/Brain2Vec",
weights_filename="vit_ba_finetuned.safetensors",
model_class=vit_base_patch8,
model_kwargs={"img_size": 128, "num_classes": 1},
device="cuda",
)
model.eval()
Brain2Vec also exists as a PyPI package and can be installed using pip:
pip install Brain2Vec
Feature extraction
One other way of using Brain2Vec is for feature extraction. We can use the Brain2Vec frozen encoder to produce 768-dimensional embeddings from T1 scans. This compact numerical representation serves as input for lightweight downstream classifiers, such as SVM or MLP models.
python inference.py brain2vec_inference --model brain2vec_base \
--img_size 128 \
--model_weights "model weights (.pth)" \
--in_dir "path to directory with niftii scans" \
--out_json "name of json file with output embeddings" \
--device cuda
python inference.py vit_inference --model vit_base \
--num_classes 1 \
--img_size 128 \
--mode regression \
--model_weights "model weights (.pth)" \
--in_dir "input directory with niftii scans" \
--out_csv "output csv file with predictions" \
--device cuda
Fine-tuning
One way of using Brain2Vec is for encoder fine-tuning. The pre-trained encoder weights can serve as initialization for Vision Transformer (ViT) models applied to downstream tasks like disease classification or brain age regression.
The only thing you need is a csv with MRID and target columns. We provide all the code needed to pre-train/fine-tune on SLURM environments on github.
srun Brain2Vec vit_finetune --img_size 128 \
--mode binary \
--model vit_base \
--label_dict "csv with MRID and target columns" \
--target "target column" \
--pretrain "pre-trained model weights (.pth)" \
--in_dir "input_dir/"
--batch_size 32 \
--out_dir "output_dir/" \
--model_name "model name (.pth)" \
--device cuda \
--lr 0.001 \
--weight_decay 0.05 \
--optimizer AdamW \
--epochs 200 \
--patience 20 \
--distributed \
--verbose
Pre-training
We trained with two A100 GPUs for 500 epochs with a patience of 20 epochs. For more training details please read the Brain2Vec paper. Using our code, you can parallelize your training on SLURM cluster environments. Other environments are supported but not tested, so bugs may appear.
srun Brain2Vec brain2vec_pretrain --img_size 128 \
--masking_type random \
--mask_ratio 0.75 \
--model brain2vec_base \
--in_dir "input_dir/"
--batch_size 32 \
--out_dir "output_dir/" \
--model_name "model name(.pth)" \
--device cuda \
--lr 0.001 \
--weight_decay 0.05 \
--optimizer AdamW \
--epochs 500 \
--patience 20 \
--distributed \
--verbose \
Usage on NiChart
You can use Brain2Vec -- for both feature extraction and disease classification/brain age regression -- on NiChart by simply dragging and dropping your NIfTI files. No GPU or installation required.
Citation
If you use our model, make sure to cite our work:
@Article{Brain2Vec2026,
author = {Spiros Maggioros and Guray Erus and Gareth Harman and George Aidinis},
journal = {arXiv:},
title = {Brain2Vec: A Self-Supervised 3D Foundation Model for Structural Brain MRI},
year = {2026},
}
- Downloads last month
- -