CytoBERT: A Foundation Model for Variable-Panel Cytometry

This repository contains the implementation of CytoBERT, a Transformer-based foundation model for single-cell cytometry data. The project supports:

• Self-supervised pretraining on large-scale CyTOF datasets
• Fine-tuning for sample-level classification

The repository also provides utilities for downloading, preprocessing, and standardizing CyTOF datasets from ImmPort.

Clone the repository. Download the cytobert-data.zip file from Zenodo and unpack it. It contains preprocessed datasets and model weights.

The code is meant to be executed in a vscode devcontainer, using the .devcontainer/devcontainer.json configuration file.

For configuring the container's environment, please customize .devcontainer/.env. Most importantly adjust the path to the extracted cytobert artifacts, and the host's CUDA version:

$ cp .devcontainer/.env.sample .devcontainer/.env
$ vim .devcontainer/.env

CYTOBERT_CUDA=${YOUR_CUDA_VERSION}
CYTOBERT_DATA=${YOUR_DATA_PATH}

Build the devcontainer and reopen the repository in the container. The devcontainer uses conda to install dependencies in the environment cytobert. The postCreateCommand will attempt to install the cytobert package into the conda environment.

Check if the setup was successful:

$ conda activate cytobert
$ python -c "import cytobert"

Please follow the notebooks for different downstream tasks:

• Tutorial_Sample_Classification.ipynb

Select cytobert as Python environment for each notebook. These tutorials demonstrate how to use pretrained CytoBERT checkpoints for different downstream tasks.

You can also run the code from a Docker Image:

wget https://zenodo.org/records/21298697/files/image.tar
docker load -i image.tar

# forward jupyter lab port, provide gpu and mount data artifacts
# make sure that the data artifacts are readable and writable to all users (chmod a+rw cytobert-data)
docker run --rm -p 8888:8888 --gpus="all" --mount ${CYTOBERT_DATA_PATH}:/workspaces/cytobert/data cytobert

You can now access localhost:8888 and run the notebooks in the notebooks folder.

.
├── .devcontainer/                         # Reproducible Docker/devcontainer environment
│
├── data_download/
│   ├── data/                              # Full Dataset manifests with meta infor and paths
│   ├── data_paths/                        # Dataset path extracted from manifests for ImmPort downloads
│   ├── Readme.md                          # Instructions for downloading CyTOF datasets
│   ├── cytof_anndata.py                   # Script for downloading CyTOF datasets
│   └── standardize_markers.ipynb          # Standardizes marker names across datasets
│
├── readfcs/                               # Modified readfcs package used for FCS parsing
│
├── src/
│   ├── cytobert/                       # Main CytoBERT package    
| 
├── notebooks/
│       ├── Tutorial_Sample_Classification     # Tutorial for sample classification
│       ├── Tutorial_Reference_Mapping         # Notebook tutorial for reference mapping
│       └── Tutorial_Cell_Classification       # Tutorial for cell classification
|
├── scripts/
│       ├── pretrain.py                        # Self-supervised pretraining pipeline
│       ├── run_pretrain.sh                    # Example pretraining shell script
│       ├── finetune.py                        # Fine-tuning and benchmarking script
│       └── run_finetune.sh                    # Shell script for main.py with different configurations
|                 
└── README.md

The complete workflow for pretraining CytoBERT is outlined below.

Activate the environment cytobert (already created during container building)

See the data_download folder for downloading and standardizing CyTOF datasets. Follow the instructions in data_download/Readme.md.

Run:

./scripts/run_pretrain.sh

This performs self-supervised pretraining on the downloaded CyTOF datasets.

Important:

• Use the provided vocabulary file:

pretrain_cytof/vocab.json

• The preprocessing statistics (bin boundaries) and checkpoints will be saved automatically.

Run:

./scripts/run_finetune.sh

Important:

• Before running the script, update the required paths inside scripts/run_finetune.sh, especially::

pretrain_run_dir/
├── checkpoints__uniform                    # checkpoint directory
└──   └── cytobert_epoch_50.pt      # pretrained CytoBERT checkpoint
├── scaler__uniform /                     # Uniform scaling (min and max value per marker)
└──   └── .npy files              # bin boundary files used during pretraining                  

label_csv/
├── SDY1535.csv
├── SDY1708.csv
├── SDY1733.csv
└── ...

• Fine-tuning on downstream datasets
• Benchmarking against baseline models
• Cross-validation experiments
• Metric generation and result saving

The script supports different:

• pooling strategies (mean, max, attention)
• dataset splits (grouped, ungrouped)
• binning schemes (unifrom, quantile)
• model configurations (cytobert, cytobert_scratch, cytoset, logistic)

• GitHub: https://github.com/bckrlab/cytobert
• Zenodo:
  • Data: https://doi.org/10.5281/zenodo.21291131
  • Docker Image: https://doi.org/10.5281/zenodo.21298697
• Docker Hub: https://hub.docker.com/repository/docker/bckrlab/cytobert