Fine-Tuning the Foundation Model
In this tutorial, we will take the pretrained model from the previous tutorial and fine-tune it for a downstream classification task. This process demonstrates how to adapt the general model to specialized classification type problems.
The associated automated script for this process is available in the scDynOmics repository under scripts/clf.ft.py.
[ ]:
import torch
import scanpy as sc
import anndata as ad, hdf5plugin
from torch.utils.data import DataLoader
from pytorch_lightning import Trainer, seed_everything
import scdynomics.utils.json as JSON
from scdynomics.utils.trainer import (
set_accelerator,
cal_accum_grad_batches,
make_callbacks,
)
from scdynomics import FT_MLP_Classifier
from scdynomics.utils.data import Multimodal_Corpus, random_split
from scdynomics.ft.pipeline import set_pretrainer, set_layers_and_embed_mod
from scdynomics import ft_kfold_classification, ft_clf_predict, ft_clf_explain
0. Configuration
First, define the paths for the downstream query dataset, the pretrained foundation model checkpoint, and hyperparameters. For this tutorial, we default to using the checkpoint generated during the pretraining tutorial.
[ ]:
# File paths
clf_adata_path = "../../data/sample/Tbx6WT.h5ad" # example data included in the scDynOmics github repository relative to this notebook
config_path = "../../data/config/tiny.adapter.clf.json" # for production, also create a modified config file, especially with higher value than 1 for trainer_hparams->max_epochs
token_dict_path = "../../data/gene_dict/token_dict.json.gz"
default_root_dir = "logs/"
# If the PT ckpt is saved from the pretraining tutorial
base_model_ckpt = "logs/lightning_logs/version_0/checkpoints/last.ckpt"
# # For randomly initialized model, set to None
# base_model_ckpt = None
# Just using the GEX modality for classification
mod1_layer = "X"
mod2_layer = None
obs_label = "cell_state"
# HVG selection parameters
nHVGs = 100
hvg_flavor = "seurat"
# Training parameters
valid_fraction = 0.1
test_fraction = 0.2
stratify_test = True
stratify_valid = True
dataloader_worker = 10
seed = 42
accelerator = "cpu"
# Apply seed
seed_everything(seed)
torch.set_float32_matmul_precision("medium")
Seed set to 42
1. Summarized Pipeline Execution
The entire fine-tuning process, including data splitting, base model loading, and classification, can be executed using the high-level :func:scdynomics.ft_kfold_classification pipeline function. This provides a unified interface for downstream adaptation.
[ ]:
# Run k-fold classification
trainer, pretrainer, classifier, clf_data = ft_kfold_classification(
config_path=config_path,
clf_adata_path=clf_adata_path,
hvg_selection_mod=None, # that automatically uses X. Specify it if not using X
nHVGs=nHVGs,
hvg_flavor=hvg_flavor,
mod1_layer=mod1_layer,
mod2_layer=mod2_layer,
obs_label=obs_label,
n_folds=1, # Just one train/valid/test split, set to >1 for k-fold CV
valid_fraction=valid_fraction,
test_fraction=test_fraction,
stratify_test=stratify_test,
stratify_valid=stratify_valid,
dataloader_worker=dataloader_worker,
base_model_ckpt=base_model_ckpt,
tuned_model_ckpt=None,
test_ckpt="best",
token_dict_path=token_dict_path,
default_root_dir=default_root_dir,
save_callback_metrics=True,
save_label_dict=True,
seed=seed,
accelerator=accelerator,
ndevices=1
)
Seed set to 42
Loaded pretrainer model
Keep modalities: [0]
Keep genes: 101; with 12 TFs
Performing single train-test split
Fold 1/1
/usr/lib/python3.11/site-packages/lightning_fabric/plugins/environments/slurm.py:204: The `srun` command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with `srun` like so: srun python /usr/lib/python3 ...
GPU available: True (cuda), used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
/usr/lib/python3.11/site-packages/pytorch_lightning/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.
| Name | Type | Params | Mode
----------------------------------------------------------
0 | ft_encoder | FT_Encoder | 5.4 M | train
1 | decision | Sequential | 13.6 K | train
2 | criterion | CrossEntropyLoss | 0 | train
3 | accuracy | MulticlassAccuracy | 0 | train
4 | f1 | MulticlassF1Score | 0 | train
5 | auroc | MulticlassAUROC | 0 | train
----------------------------------------------------------
2.2 M Trainable params
3.2 M Non-trainable params
5.4 M Total params
21.762 Total estimated model params size (MB)
294 Modules in train mode
0 Modules in eval mode
Initializing new FT model
`Trainer.fit` stopped: `max_epochs=1` reached.
Restoring states from the checkpoint path at logs/fold_1/lightning_logs/version_0/checkpoints/epoch=0-step=34.ckpt
Loaded model weights from the checkpoint at logs/fold_1/lightning_logs/version_0/checkpoints/epoch=0-step=34.ckpt
Testing the model
/usr/lib/python3.11/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: No negative samples in targets, false positive value should be meaningless. Returning zero tensor in false positive score
warnings.warn(*args, **kwargs) # noqa: B028
/usr/lib/python3.11/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: No positive samples in targets, true positive value should be meaningless. Returning zero tensor in true positive score
warnings.warn(*args, **kwargs) # noqa: B028
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ DataLoader 0 ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ test.accuracy │ 0.5526315569877625 │ │ test.auroc │ 0.4736842215061188 │ │ test.f1 │ 0.38847118616104126 │ │ test.loss │ 0.7271869778633118 │ └───────────────────────────┴───────────────────────────┘
Fold 1/1 done
Fold 1 metrics: {'train.loss': 0.5477855205535889, 'train.grad_norm': 1.4536314010620117, 'vali.loss': 0.7176894545555115, 'vali.accuracy': 0.5625, 'vali.f1': 0.34761905670166016, 'vali.auroc': 0.6041666269302368, 'test.loss': 0.7271869778633118, 'test.accuracy': 0.5526315569877625, 'test.f1': 0.38847118616104126, 'test.auroc': 0.4736842215061188}
2. Detailed Step-by-Step Breakdown
Below, we deconstruct the :func:scdynomics.ft_kfold_classification pipeline to explore the individual steps involved in data preprocessing, model initialization, and fine-tuning execution.
2.1 Data Preprocessing
Load the query dataset meant for fine-tuning. We expect the target labels for the classification task (e.g., cell types or states) to be stored within the .obs attribute of the AnnData object.
2.1.1 Highly Variable Gene (HVG) Selection
(Optional) We can subset the dataset to highly variable genes to reduce dimensionality and focus the classifier on the most informative features.
[4]:
adata = ad.read_h5ad(clf_adata_path)
sc.pp.highly_variable_genes(
adata,
n_top_genes=nHVGs,
subset=True,
flavor= hvg_flavor,
)
2.1.2 Load Pretrained Model and Align Feature Space
Here, we load the :class:scdynomics.Masked_Pretrainer base model and automatically align its feature space (token embeddings) with the query dataset using the :func:scdynomics.ft.pipeline.set_pretrainer utility.
[5]:
# Set the layers
layers, embed_mod = set_layers_and_embed_mod(mod1_layer, mod2_layer)
# Load the pretrained base model and the corresponding data corpus for classification
pretrainer, clf_data = set_pretrainer(
clf_data = Multimodal_Corpus(
adata = adata,
label_key = obs_label,
layers = layers,
backed = False,
),
embed_mod = embed_mod,
base_model_ckpt = base_model_ckpt,
token_dict_path = token_dict_path,
)
Loaded pretrainer model
Keep modalities: [0]
Keep genes: 101; with 12 TFs
2.1.3 Dataset Splitting and Dataloader Initialization
We partition the aligned corpus into training, validation, and test sets. These are then wrapped in PyTorch :class:torch.utils.data.DataLoader objects for efficient batch processing.
[6]:
# Load the config file
config = JSON.decode(config_path)
dl_hparams = config["dataloader_hparams"]
# Split train and test data
train_list, valid_list, test_list = random_split(
clf_data,
test_fraction=test_fraction,
valid_fraction=valid_fraction,
stratified_test=stratify_test,
stratified_valid=stratify_valid,
random_seed=seed,
)
# Create dataloaders
train_data = DataLoader(
train_list[0],
num_workers = dataloader_worker,
batch_size = dl_hparams["batch_per_device"],
shuffle = True
)
val_data = DataLoader(
valid_list[0],
num_workers = dataloader_worker,
batch_size = dl_hparams["batch_per_device"],
shuffle = False
) if valid_list[0] is not None else None
test_data = DataLoader(
test_list[0],
num_workers = dataloader_worker,
batch_size = dl_hparams["batch_per_device"],
shuffle = False
) if test_list[0] is not None else None
2.2 Model Initialization
We initialize the :class:scdynomics.FT_MLP_Classifier. This downstream module encapsulates our pretrained base model, freezing its core weights while injecting Adapter modules for parameter efficient fine-tuning and appending an adaptable Multi-Layer Perceptron (MLP) head for classification.
[7]:
classifier = FT_MLP_Classifier(
base_model = pretrainer,
n_classes = len(clf_data.label_dict),
steps_per_epoch = len(train_data),
**config["classifier_hparams"],
)
2.3 Trainer Setup & Execution
We configure the PyTorch Lightning :class:lightning.pytorch.Trainer and invoke the .fit() method to begin the supervised fine-tuning process on the downstream task.
[ ]:
# Init the Trainer object
trainer = Trainer(
devices = 1,
accelerator = accelerator,
default_root_dir = default_root_dir + f"fold_1/",
accumulate_grad_batches = cal_accum_grad_batches(config, 1),
callbacks = make_callbacks(config),
**config["trainer_hparams"]
)
# Train the model
trainer.fit(
model=classifier,
train_dataloaders=train_data,
val_dataloaders=val_data,
)
# Get the callback metrics
callback_metrics = {
k: float(v)
for k,v in trainer.callback_metrics.items()
}
/usr/lib/python3.11/site-packages/lightning_fabric/plugins/environments/slurm.py:204: The `srun` command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with `srun` like so: srun python /usr/lib/python3 ...
GPU available: True (cuda), used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
/usr/lib/python3.11/site-packages/pytorch_lightning/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.
| Name | Type | Params | Mode
----------------------------------------------------------
0 | ft_encoder | FT_Encoder | 5.4 M | train
1 | decision | Sequential | 13.6 K | train
2 | criterion | CrossEntropyLoss | 0 | train
3 | accuracy | MulticlassAccuracy | 0 | train
4 | f1 | MulticlassF1Score | 0 | train
5 | auroc | MulticlassAUROC | 0 | train
----------------------------------------------------------
2.2 M Trainable params
3.2 M Non-trainable params
5.4 M Total params
21.762 Total estimated model params size (MB)
294 Modules in train mode
0 Modules in eval mode
`Trainer.fit` stopped: `max_epochs=1` reached.
2.4 Test Set Evaluation
After training, we evaluate the best model checkpoint against the held-out test set to calculate the final accuracy, F1 score, and AUROC metrics.
[9]:
# test the model
trainer.test(
model=classifier,
dataloaders=test_data,
ckpt_path="best"
)
# Get the callback metrics
for k,v in trainer.callback_metrics.items():
callback_metrics[k] = float(v)
print(f"callback_metrics: {callback_metrics}")
Restoring states from the checkpoint path at logs/fold_1/lightning_logs/version_0/checkpoints/epoch=0-step=34.ckpt
Loaded model weights from the checkpoint at logs/fold_1/lightning_logs/version_0/checkpoints/epoch=0-step=34.ckpt
/usr/lib/python3.11/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: No negative samples in targets, false positive value should be meaningless. Returning zero tensor in false positive score
warnings.warn(*args, **kwargs) # noqa: B028
/usr/lib/python3.11/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: No positive samples in targets, true positive value should be meaningless. Returning zero tensor in true positive score
warnings.warn(*args, **kwargs) # noqa: B028
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Test metric ┃ DataLoader 0 ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩ │ test.accuracy │ 0.5526315569877625 │ │ test.auroc │ 0.3815789520740509 │ │ test.f1 │ 0.38847118616104126 │ │ test.loss │ 0.6895833015441895 │ └───────────────────────────┴───────────────────────────┘
callback_metrics: {'train.loss': 0.9454113245010376, 'train.grad_norm': 1.589394211769104, 'vali.loss': 0.6885690093040466, 'vali.accuracy': 0.5625, 'vali.f1': 0.34761905670166016, 'vali.auroc': 0.375, 'test.loss': 0.6895833015441895, 'test.accuracy': 0.5526315569877625, 'test.f1': 0.38847118616104126, 'test.auroc': 0.3815789520740509}
3. Apply Fine-Tuned Classifier for Prediction
Once fine-tuned, the classifier can be used to predict labels for new cells. The :func:scdynomics.ft_clf_predict function automates the inference process and generates a classification report if true labels are provided.
The associated automated script is available in the scDynOmics repository under scripts/clf.predict.py.
[ ]:
# Set the paths
label_dict_path = default_root_dir + f"label_dict.json"
tuned_model_ckpt = default_root_dir + f"fold_1/lightning_logs/version_0/checkpoints/epoch=0-step=34.ckpt"
# Run prediction
all_preds, clf_report = ft_clf_predict(
clf_adata=clf_data.adata, # Directly use the adata since it's already preprocessed and aligned
mod1_layer=mod1_layer,
mod2_layer=mod2_layer,
obs_label=obs_label,
label_dict_path=label_dict_path,
token_dict_path=token_dict_path,
base_model_ckpt=base_model_ckpt,
tuned_model_ckpt=tuned_model_ckpt,
outpath= default_root_dir + f"preds.json",
accelerator = accelerator,
seed=seed,
)
Seed set to 42
Loaded pretrainer model
Keep modalities: [0]
Keep genes: 101; with 12 TFs
Predictions saved to logs/preds.json
Classification report saved to logs/preds.report.json
/usr/lib/python3.11/site-packages/sklearn/metrics/_classification.py:1833: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.
_warn_prf(average, modifier, f"{metric.capitalize()} is", result.shape[0])
/usr/lib/python3.11/site-packages/sklearn/metrics/_classification.py:1833: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.
_warn_prf(average, modifier, f"{metric.capitalize()} is", result.shape[0])
/usr/lib/python3.11/site-packages/sklearn/metrics/_classification.py:1833: UndefinedMetricWarning: Precision is ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.
_warn_prf(average, modifier, f"{metric.capitalize()} is", result.shape[0])
4. Explain Fine-Tuned Classifier
To ensure interpretability, scDynOmics includes an explanation module utilizing Integrated Gradients. This allows us to quantify the contribution of individual features (genes) to the model’s decisions.
The associated automated script is available in the scDynOmics repository under scripts/clf.explain.py.
[ ]:
# Explanation settings
explain_sample_limit = 2
sample_method = 'top'
explain_batch_size = 1
explain_step = 100
# Run the explanation
explain_result = ft_clf_explain(
clf_adata = clf_data.adata, # Directly use the adata since it's already preprocessed and aligned
mod1_layer = mod1_layer,
mod2_layer = mod2_layer,
dataloader_worker = dataloader_worker,
token_dict_path = token_dict_path,
base_model_ckpt = base_model_ckpt,
tuned_model_ckpt = tuned_model_ckpt,
explain_sample_limit = explain_sample_limit,
sample_method= sample_method,
explain_batch_size = explain_batch_size,
explain_step = explain_step,
outpath = default_root_dir + f"explain.txt",
accelerator = accelerator,
seed = seed,
)
Seed set to 42
Loaded pretrainer model
Keep modalities: [0]
Keep genes: 101; with 12 TFs
Launching Class 0 Explanation...
Explainer: 0 samples predicted as class 0.
Class 0 Explanation Done: 0.09 minutes
Launching Class 1 Explanation...
Explainer: 190 samples predicted as class 1.
Processing 1th batch...
Finished batch 1: 0.11 minutes
Processing 2th batch...
Finished batch 2: 0.10 minutes
Class 1 Explanation Done: 0.30 minutes
Explanation saved to logs/explain.txt
4. TensorBoard Visualization
The training progress and metrics can be monitored using TensorBoard by running the following command in the terminal:
tensorboard --logdir logs/fold_1
5. Cleanup
Remove the generated logs and checkpoints to free up space.
[ ]:
import shutil
shutil.rmtree(default_root_dir)