Spatium Fine-tuning Training Example

spaProFormer supports fine-tuning on multiple downstream tasks for spatial proteomics analysis. Users can configure the downstream objective through the task parameter.

Selecting downstream tasks

The fine-tuning task is specified in the configuration dictionary:

"task": "cell_type_prediction"

Supported tasks include:

  • cell_type_prediction

    Predict cell types from spatial proteomics profiles.

  • Prototype_classification

    Perform prototype-based classification.

  • neighborhood_identify

    Predict spatial neighborhood composition.

  • panel_expansion_continuous_new

    Perform continuous protein panel expansion and imputation.

  • image_integration

    Integrate spatial proteomics data with image features.

  • reconstruction

    Reconstruct masked protein expression profiles.

  • label_transfer

    Transfer annotations from reference datasets.

Dataset input

The input data should be stored as Zarr format and provided through zarr_path.

A single dataset can be provided:

zarr_path = "path/to/dataset.zarr"

For predefined training and validation datasets, provide multiple Zarr paths:

zarr_path = [
    "path/to/train.zarr",
    "path/to/validation.zarr"
]

When multiple files are provided, enable file-based splitting:

split_by_file = True

For a single dataset with internal splitting:

split_by_file = False

Complete fine-tuning example

The complete training script is shown below.

This example demonstrates:

  • loading a pretrained spaProFormer model

  • configuring downstream tasks

  • preparing Zarr datasets

  • initializing the fine-tuning model

  • training with PyTorch Lightning

  1from model import spaProFormer
  2from fine_tune import spaProFormerFinetune
  3from dataloader import MerlinDataModule
  4import pytorch_lightning as pl
  5from pytorch_lightning.loggers import WandbLogger
  6from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
  7import pandas as pd
  8import yaml
  9import os
 10import torch
 11import zarr
 12import numpy as np
 13os.environ["WANDB_MODE"] = "offline"
 14
 15if __name__ == '__main__':
 16    config = {
 17        'pretrained_path': "/path/to/final.ckpt", 
 18        'retake_training': False,
 19        'dim_feedforward': 512,
 20        'nheads': 16,
 21        'masking_p': 0.15,
 22        'nlayers': 12,
 23        'dropout': 0.2,
 24        'dim_model': 256,
 25        'batch_first': True,
 26        'n_tokens': 278, 
 27        'batch_size': 128,
 28        # 'batch_size': 64,
 29        'context_length': 239,
 30        'lr': 5e-5,
 31        'weight_decay': 0.1,
 32        'warmup': 50000, 
 33        'max_epochs': 20, 
 34        'task': 'cell_type_prediction',
 35        # 'task': 'Prototype_classification',
 36        # 'task': 'neighborhood_identify',
 37        # 'task': 'panel_expansion_continuous_new',
 38        # 'task': 'image_integration',
 39        # 'task': 'reconstruction',
 40        # 'task': 'label_transfer',
 41        'finetune_mode': 'full',
 42        'supervised_task': False,
 43        'learnable_pe': True,
 44        }
 45
 46
 47    save_ckpt_dir = '/path/to/save/ckpt'
 48    os.makedirs(save_ckpt_dir, exist_ok=True)
 49
 50    with open(os.path.join(save_ckpt_dir, "train_config.yaml"), "w") as f:
 51        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 52
 53    pl.seed_everything(42)
 54
 55
 56    # zarr_path = '/path/to/input'
 57    zarr_path = [
 58        "path/to/train_dataset.zarr",
 59        "path/to/validation_dataset.zarr",
 60    ]
 61 
 62    if isinstance(zarr_path, list):
 63        z = zarr.open(zarr_path[0], mode='r')
 64    else:
 65        z = zarr.open(zarr_path, mode='r')
 66    if config['task'] == 'neighborhood_identify':
 67        z = zarr.open(zarr_path, mode='r')
 68        if 'neighbor_ratio' in z and z.get('neighbor_ratio') is not None:
 69            num_cell_types = np.array(z.get('neighbor_ratio')).shape[1]
 70    elif config['task'] == 'label_transfer':
 71        z = zarr.open(zarr_path[0], mode='r')
 72        labels = np.array(z.get('label'))
 73        num_cell_types = labels.max() + 1
 74    elif 'label' in z and z.get('label') is not None:
 75        z = zarr.open(zarr_path, mode='r')
 76        labels = np.array(z.get('label'))
 77        num_cell_types = labels.max() + 1
 78    else:
 79        num_cell_types = None
 80
 81    dm = MerlinDataModule(
 82        zarr_path=zarr_path,
 83        batch_size=config['batch_size'],
 84        task=config['task'],
 85        split_by_file=True,
 86        context_length=config['context_length'],
 87    )
 88    dm.setup()
 89
 90
 91    pretrained_model = spaProFormer.load_from_checkpoint(config['pretrained_path'])
 92
 93    if config['task'] == 'panel_expansion_continuous':
 94        continuous_dim = zarr.open(zarr_path, mode='r').get('continuous').shape[1]
 95    elif config['task'] == 'panel_expansion_continuous_new':
 96        continuous_dim = zarr.open(zarr_path[0], mode='r').get('continuous').shape[1]
 97    else:
 98        continuous_dim = None
 99    fine_tune_model = spaProFormerFinetune(
100        pretrained_model=pretrained_model,
101        num_cell_types=int(num_cell_types) if num_cell_types is not None else None,
102        drop_out=config['dropout'],
103        task=config['task'],
104        lr=1e-5,
105        finetune_mode=config['finetune_mode'],
106        continuous_dim=continuous_dim,
107        graph_pe_dim=8
108    )
109
110    wandb_logger = WandbLogger(
111        name=f"your wandb logger name",
112        project="your wandb logger project",
113        entity="your wandb logger entity",
114    )
115
116    checkpoint_callback = ModelCheckpoint(
117        dirpath=save_ckpt_dir,
118        filename="scProFormerFineTune-{epoch:02d}-{val_accs:.4f}",
119        save_top_k=-1,
120        every_n_epochs=1,
121        save_on_train_epoch_end=True,
122    )
123    lr_monitor = LearningRateMonitor(logging_interval="step")
124
125    trainer = pl.Trainer(
126        accelerator="gpu",
127        devices=-1,
128        max_epochs=config['max_epochs'],
129        logger=wandb_logger,
130        strategy="ddp_find_unused_parameters_true",
131        callbacks=[checkpoint_callback, lr_monitor],
132        log_every_n_steps=5,
133        precision="16-mixed",
134        use_distributed_sampler=False if config['task'] == 'pancancer_engine' else True,
135        num_sanity_val_steps=0,
136
137    )
138
139
140    trainer.fit(fine_tune_model, datamodule=dm)