Pretraining
In this tutorial, we will pretrain a masked transformer model on our single-cell multiomics data using scDynOmics. The pretraining process involves masking a portion of the input data and training the model to predict the masked values, which helps the model learn meaningful representations of the data.
The associated script is available in the scDynOmics repository under scripts/pt.py.
[ ]:
import torch
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.utils.data import Multimodal_Corpus
from scdynomics.pt.reconstructor import load_or_init
from scdynomics.pt.pipeline import bimodal as bimodal_pt
0. Configuration
For the purpose of this tutorial, the configuration parameters are explicitly defined below to make them easy to explore and modify.
Note: In a production environment, especially when using SLURM or other high-performance computing (HPC) resource distribution systems, these configurations are typically parsed from command-line arguments. This allows for greater flexibility and seamless integration with job scheduling pipelines.
[2]:
# File paths
data_corpus_path = "../../data/sample/mm10.100sample.h5ad" # example data included in the scDynOmics repository relative to this notebook
config_path = "../../data/config/mm10.pt.sample.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"
pretrainer_ckpt = None # Set to a path if resuming from a checkpoint
# Data and layers
mod1_layer = "X"
mod2_layer = "pan_promoter"
# Training parameters
valid_fraction = 0.1
dataloader_worker = 10
seed = 42
# Apply seed
seed_everything(seed)
torch.set_float32_matmul_precision("medium")
Seed set to 42
1. Summarized Pipeline Execution
You can run the entire pretraining process using the high-level :func:scdynomics.pt.pipeline.bimodal function. This function abstracts away the individual data loading and training steps.
[ ]:
# Run pretraining
trainer, pretrainer_model = bimodal_pt(
config_path = config_path,
data_corpus_path = data_corpus_path,
mod1_layer = mod1_layer,
mod2_layer = mod2_layer,
valid_fraction = valid_fraction,
dataloader_worker = dataloader_worker,
pretrainer_ckpt = pretrainer_ckpt,
token_dict_path = token_dict_path,
default_root_dir = default_root_dir,
seed = seed,
ndevices = 1,
accelerator = 'cpu',
)
Seed set to 42
/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 | vbinning | Value_Binning | 0 | train
1 | input_dropout | Input_Dropout | 0 | train
2 | embedder | Genome_Embedder | 2.2 M | train
3 | encoder | Encoder | 75.8 M | train
4 | pool | Sequential | 0 | train
5 | mod0_rebuilder | Linear | 1.5 K | train
6 | mod1_rebuilder | Linear | 1.5 K | train
7 | criterion | CrossEntropyLoss | 0 | train
------------------------------------------------------------
77.9 M Trainable params
0 Non-trainable params
77.9 M Total params
311.676 Total estimated model params size (MB)
185 Modules in train mode
0 Modules in eval mode
Random initialized pretrainer model
`Trainer.fit` stopped: `max_epochs=1` reached.
2. Detailed Step-by-Step Breakdown
Below, we break down what happens under the hood of the :func:scdynomics.pt.pipeline.bimodal function to give you a deeper understanding of the framework’s mechanics.
2.1 Data Loading
We load our data into a :class:scdynomics.utils.data.Multimodal_Corpus and split it into training and validation sets. Then, we wrap them in PyTorch DataLoaders for efficient batching.
[4]:
# Load the config file
config = JSON.decode(config_path)
dl_hparams = config["dataloader_hparams"]
# Load training data corpus
full_data = Multimodal_Corpus(
data_corpus_path,
layers = [mod1_layer, mod2_layer],
backed = False,
)
# Split train and validation data
train_data_size = int(1-valid_fraction * len(full_data))
train_data, val_data = torch.utils.data.random_split(
full_data,
[train_data_size, len(full_data) - train_data_size]
)
train_data = DataLoader(
train_data,
num_workers = dataloader_worker,
batch_size = dl_hparams["batch_per_device"],
shuffle = True
)
val_data = DataLoader(
val_data,
num_workers = dataloader_worker,
batch_size = dl_hparams["batch_per_device"],
shuffle = False
)
2.2 Model Initialization
We initialize our :class:scdynomics.Masked_Pretrainer. If a checkpoint path is provided, we load the model from the checkpoint to resume training.
[5]:
# Load or initialize the pretrainer model
pretrainer_model = load_or_init(
ckpt_path = pretrainer_ckpt,
token_dict_path = token_dict_path,
config = config,
)
Random initialized pretrainer model
2.3 Trainer Setup & Execution
Finally, we set up PyTorch Lightning’s Trainer with our callbacks, like Model Checkpointing, and start the training process
[6]:
# Init trainer
trainer = Trainer(
devices = 1,
accelerator = "cpu", # or "cuda",
default_root_dir = default_root_dir,
accumulate_grad_batches = cal_accum_grad_batches(config, 1),
callbacks = make_callbacks(config),
**config["trainer_hparams"]
)
# Train the model
trainer.fit(
model = pretrainer_model,
train_dataloaders = train_data,
val_dataloaders = val_data,
ckpt_path = pretrainer_ckpt,
)
/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 | vbinning | Value_Binning | 0 | train
1 | input_dropout | Input_Dropout | 0 | train
2 | embedder | Genome_Embedder | 2.2 M | train
3 | encoder | Encoder | 75.8 M | train
4 | pool | Sequential | 0 | train
5 | mod0_rebuilder | Linear | 1.5 K | train
6 | mod1_rebuilder | Linear | 1.5 K | train
7 | criterion | CrossEntropyLoss | 0 | train
------------------------------------------------------------
77.9 M Trainable params
0 Non-trainable params
77.9 M Total params
311.676 Total estimated model params size (MB)
185 Modules in train mode
0 Modules in eval mode
`Trainer.fit` stopped: `max_epochs=1` reached.
3. TensorBoard Visualization
The training progress and metrics can be monitored using TensorBoard by running the following command in the terminal:
tensorboard --logdir logs
4. Cleanup
WARNING: Don’t run this cell if you want to load the checkpoint for the fine-tuning tutorial.
[ ]:
import shutil
shutil.rmtree(default_root_dir)