diff --git a/.gitignore b/.gitignore index 5187d86..9097fb2 100644 --- a/.gitignore +++ b/.gitignore @@ -180,3 +180,10 @@ figures/ # hydra .hydra/ new_env/ + +data/ +checkpoints/ +scripts/ +examples/ +outputs/ +conf/ diff --git a/README.md b/README.md index 26a20ce..fe2e646 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,7 @@ uv pip install -e . ### Requirements Transcriptformer has the following core dependencies: -- PyTorch (<=2.5.1, as 2.6.0+ may cause pickle errors) +- PyTorch (>=2.5.1, <2.13) # original says PyTorch (<=2.5.1, as 2.6.0+ may cause pickle errors) - PyTorch Lightning - anndata - scanpy diff --git a/conf/train_config.yaml b/conf/train_config.yaml new file mode 100644 index 0000000..3f24dba --- /dev/null +++ b/conf/train_config.yaml @@ -0,0 +1,88 @@ +defaults: + - _self_ # Load the current config file + +train: + # === Checkpointing & Output === + checkpoint_dir: null # Path to source checkpoint directory (required) + resume_artifact_dir: null # Path to previous training output for resuming (optional) + resume_mode: weights # 'weights' (load weights only) or 'lightning' (resume full state) + output_dir: ./train_outputs/tf_assay_expanded # Output directory for new artifacts + + checkpoint: + save_top_k: 1 # Number of best checkpoints to keep (by val/total_loss) + save_last: true # Always save the last checkpoint + + live_metrics: + enabled: true # Export per-epoch summaries and live plots while training + every_n_epochs: 1 # Refresh metrics/plots every N epochs + + # === Data === + train_files: [] # List of training .h5ad files + val_files: [] # List of validation .h5ad files + expanded_assay_vocab: null # Path to expanded assay vocab JSON (optional) + obs_assay_col: assay # Required AnnData.obs column used to create aux tokens for model forward + obs_keys: null # Optional list of extra AnnData.obs fields to carry in batch.obs (metadata only) + + data_config: + _target_: transcriptformer.data.dataclasses.DataConfig + gene_col_name: ensembl_id # Column name in AnnData.var containing gene identifiers + clip_counts: 30.0 # Maximum count value (higher values will be clipped), Null means inherit from checkpoint model config + filter_to_vocabs: true # Whether to filter genes to only those in the vocabulary, Null means inherit from checkpoint model config + filter_outliers: 3.0 # Standard deviation threshold for filtering outlier cells (0.0 = no filtering), Null means inherit from checkpoint model config + normalize_to_scale: 0.0 # Scale factor for count normalization (0 = no normalization) + sort_genes: null # Whether to sort genes by expression level, Null means inherit from checkpoint model config + randomize_genes: true # Whether to randomize gene order, Null means inherit from checkpoint model config + min_expressed_genes: 200 # Minimum number of expressed genes required per cell, Null means inherit from checkpoint model config + use_raw: "auto" # Whether to use .raw.X (True), .X (False), or auto-detect (auto/null), Null means inherit from checkpoint model config + remove_duplicate_genes: false # Remove duplicate genes instead of raising + n_data_workers: 4 # Data workers used by data module config + + # === Dataloader === + batch_size: 8 # Batch size + num_workers: 4 # DataLoader workers + use_oom_dataloader: true # Use backed OOM dataset + file-aware batching (lower memory, mostly file-local batches) + enable_file_aware_batching: true # Training-only and DDP-safe: keeps file-local OOM batches; set false to use Lightning DistributedSampler path + oom_batches_per_file: 1 # In OOM mode, number of consecutive batches drawn from one file before interleaving + + + # === Training Loop === + # Recommended: Aim for a global batch size of 4–5 million tokens (examples * seq_len * num_gpus * num_nodes * accumulate_grad_batches) + max_epochs: 5 # Number of epochs + precision: 16-mixed # Precision (16-mixed, 32, bf16, etc.) + device: auto # Preferred device policy ('auto', 'cpu', 'cuda', 'mps'), same resolution path as inference + num_gpus: 1 # Requested GPU count (1, >1, or -1 for all available on each node) + num_nodes: 1 # Number of nodes (for DDP) + log_every_n_steps: 10 # Logging frequency + seed: 42 # Random seed + gradient_clip_val: 1.0 # Max norm for gradient clipping (improves stability) + accumulate_grad_batches: 6 # Number of steps to accumulate gradients before optimizer step + + # === Optimizer & Scheduler === + lr: 5.5e-5 # Learning rate, try 1.5e-4 if train from scratch + weight_decay: 0.05 # AdamW weight decay, try 0.1 if train from scratch + adam_beta1: 0.9 # AdamW beta1 + adam_beta2: 0.95 # AdamW beta2 + adam_eps: 1e-8 # AdamW epsilon + warmup_ratio: 0.1 # LR warmup ratio + min_lr_ratio: 0.18 # Minimum LR as fraction of initial, cap at 1e-5 + + # === Loss Weights === + loss_config: + _target_: transcriptformer.data.dataclasses.LossConfig + gene_id_loss_weight: 1.0 # Gene ID prediction loss weight + softplus_approx: true # Use softplus approximation in loss path + + # === Assay Embedding Init === + init_default_source: unknown # Default source for new assay tokens + assay_init_map: [] # List of 'new=source' mappings + + # === Freezing & Training Policy === + freeze_transformer: false # Freeze transformer encoder + unfreeze_last_n_transformer_blocks: 0 # If freeze_transformer=true, unfreeze only the last N encoder blocks (0 keeps full encoder frozen) + freeze_gene_embeddings: false # Freeze gene embeddings + freeze_count_head: false # Freeze count decoder + freeze_gene_head: false # Freeze gene decoder + train_aux_only: false # Train only aux (assay) embeddings + + # === Data Preprocessing === + shuffle_expressed_each_batch: false # Shuffle expressed genes per batch diff --git a/pyproject.toml b/pyproject.toml index 533b13c..a424745 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ "hydra-core==1.3.2", "numpy==2.2.6", "pandas==2.2.2", - "psutil==5.9", + "psutil>=5.9,<8", "pynvml==12", "pytest==8.4", "pytorch-lightning==2.5.1", @@ -31,7 +31,7 @@ dependencies = [ "scanpy==1.11.2", "scipy==1.15.3", "timeout-decorator==0.5", - "torch==2.5.1", + "torch>=2.5.1,<2.13", ] optional-dependencies.build = [ diff --git a/src/transcriptformer/cli/__init__.py b/src/transcriptformer/cli/__init__.py index 362229f..9327ee7 100644 --- a/src/transcriptformer/cli/__init__.py +++ b/src/transcriptformer/cli/__init__.py @@ -69,6 +69,7 @@ from omegaconf import OmegaConf from transcriptformer.model.inference import run_inference +from transcriptformer.train.engine import run_train_from_dict, setup_runtime_for_training # Suppress annoying warnings warnings.filterwarnings("ignore", category=FutureWarning, module="anndata") @@ -284,6 +285,102 @@ def setup_download_data_parser(subparsers): ) +def setup_train_parser(subparsers): + """Setup the parser for the train command.""" + parser = subparsers.add_parser( + "train", + help="Train or continue training with expanded assay vocab", + description="Fine-tune TranscriptFormer with expanded assay tokens and optional freezing.", + ) + + parser.add_argument("--checkpoint-dir", required=True, help="Base artifact directory with config.json/model_weights.pt") + parser.add_argument("--output-dir", required=True, help="Output artifact directory") + parser.add_argument("--train-file", action="append", required=True, help="Training .h5ad file (repeatable)") + parser.add_argument("--val-file", action="append", default=[], help="Validation .h5ad file (repeatable)") + + parser.add_argument("--expanded-assay-vocab", help="Expanded assay_vocab.json path") + parser.add_argument("--obs-assay-col", default="assay", help="obs assay column") + parser.add_argument("--gene-col-name", default="ensembl_id", help="adata.var gene ID column") + parser.add_argument("--filter-to-vocabs", action="store_true", default=True, help="filter genes to vocabulary") + parser.add_argument("--filter-outliers", type=float, default=0.0, help="cell outlier filtering threshold") + parser.add_argument("--sort-genes", action="store_true", help="sort genes by expression") + parser.add_argument("--randomize-genes", action="store_true", help="randomize gene order") + parser.add_argument("--min-expressed-genes", type=int, default=0, help="minimum expressed genes per cell") + parser.add_argument("--n-data-workers", type=int, default=4, help="DataConfig n_data_workers value") + + parser.add_argument("--resume-artifact-dir", default=None, help="Resume from previous output artifact directory") + parser.add_argument( + "--resume-mode", + default="weights", + choices=["weights", "lightning"], + help="Resume mode: weights (supports runtime policy changes) or lightning (resume optimizer/scheduler state)", + ) + + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument("--max-epochs", type=int, default=5) + parser.add_argument("--precision", default="16-mixed") + parser.add_argument( + "--num-gpus", + type=int, + default=1, + help="Number of GPUs per node to use for training (1=single, -1=all available, >1=specific number)", + ) + parser.add_argument("--num-nodes", type=int, default=1) + parser.add_argument( + "--device", + default="auto", + choices=["auto", "cpu", "cuda", "mps"], + help="Preferred device policy for training", + ) + + parser.add_argument("--lr", type=float, default=5.5e-5) + parser.add_argument("--weight-decay", type=float, default=0.05) + parser.add_argument("--adam-beta1", type=float, default=0.9) + parser.add_argument("--adam-beta2", type=float, default=0.95) + parser.add_argument("--adam-eps", type=float, default=1e-8) + parser.add_argument("--warmup-ratio", type=float, default=0.1) + parser.add_argument("--min-lr-ratio", type=float, default=0.1) + + parser.add_argument("--gene-id-loss-weight", type=float, default=1.0) + parser.add_argument("--softplus-approx", action=argparse.BooleanOptionalAction, default=True) + + parser.add_argument("--init-default-source", default="unknown") + parser.add_argument("--assay-init-map", action="append", default=[], help="new_assay=source_assay mapping") + + parser.add_argument("--freeze-transformer", action="store_true") + parser.add_argument( + "--unfreeze-last-n-transformer-blocks", + type=int, + default=0, + help="When --freeze-transformer is set, unfreeze only the last N encoder blocks (0 keeps full transformer frozen)", + ) + parser.add_argument("--freeze-gene-embeddings", action="store_true") + parser.add_argument("--freeze-count-head", action="store_true") + parser.add_argument("--freeze-gene-head", action="store_true") + parser.add_argument("--train-aux-only", action="store_true") + + parser.add_argument("--shuffle-expressed-each-batch", action="store_true") + parser.add_argument("--clip-counts", type=float, default=30.0) + parser.add_argument("--normalize-to-scale", type=float, default=0.0) + parser.add_argument("--use-raw", action="store_true") + parser.add_argument("--remove-duplicate-genes", action="store_true") + parser.add_argument("--use-oom-dataloader", action="store_true") + parser.add_argument( + "--file-aware-batching", + action=argparse.BooleanOptionalAction, + default=True, + help="Keep OOM batches mostly file-local; disable to fall back to Lightning DistributedSampler batching", + ) + parser.add_argument( + "--oom-batches-per-file", + type=int, + default=1, + help="In OOM mode, number of consecutive batches to draw from one file before interleaving", + ) + parser.add_argument("--seed", type=int, default=42) + + def run_inference_cli(args): """Run inference using command line arguments.""" # Only print logo if not in distributed mode (avoids duplicates) @@ -452,6 +549,68 @@ def run_download_data_cli(args): sys.exit(1) +def run_train_cli(args): + """Run training from CLI args.""" + setup_runtime_for_training() + cfg = { + "checkpoint_dir": args.checkpoint_dir, + "resume_artifact_dir": args.resume_artifact_dir, + "resume_mode": args.resume_mode, + "output_dir": args.output_dir, + "train_files": args.train_file, + "val_files": args.val_file, + "expanded_assay_vocab": args.expanded_assay_vocab, + "obs_assay_col": args.obs_assay_col, + "data_config": { + "gene_col_name": args.gene_col_name, + "clip_counts": args.clip_counts, + "filter_to_vocabs": args.filter_to_vocabs, + "filter_outliers": args.filter_outliers, + "normalize_to_scale": args.normalize_to_scale, + "sort_genes": args.sort_genes, + "randomize_genes": args.randomize_genes, + "min_expressed_genes": args.min_expressed_genes, + "use_raw": args.use_raw, + "remove_duplicate_genes": args.remove_duplicate_genes, + "n_data_workers": args.n_data_workers, + }, + "batch_size": args.batch_size, + "num_workers": args.num_workers, + "max_epochs": args.max_epochs, + "precision": args.precision, + "device": args.device, + "num_gpus": args.num_gpus, + "num_nodes": args.num_nodes, + "lr": args.lr, + "weight_decay": args.weight_decay, + "adam_beta1": args.adam_beta1, + "adam_beta2": args.adam_beta2, + "adam_eps": args.adam_eps, + "warmup_ratio": args.warmup_ratio, + "min_lr_ratio": args.min_lr_ratio, + "loss_config": { + "gene_id_loss_weight": args.gene_id_loss_weight, + "softplus_approx": args.softplus_approx, + }, + "init_default_source": args.init_default_source, + "assay_init_map": args.assay_init_map, + "freeze_transformer": args.freeze_transformer, + "unfreeze_last_n_transformer_blocks": args.unfreeze_last_n_transformer_blocks, + "freeze_gene_embeddings": args.freeze_gene_embeddings, + "freeze_count_head": args.freeze_count_head, + "freeze_gene_head": args.freeze_gene_head, + "train_aux_only": args.train_aux_only, + "shuffle_expressed_each_batch": args.shuffle_expressed_each_batch, + "use_oom_dataloader": args.use_oom_dataloader, + "enable_file_aware_batching": args.file_aware_batching, + "oom_batches_per_file": args.oom_batches_per_file, + "seed": args.seed, + } + + result = run_train_from_dict(cfg) + print(f"Training complete. Artifacts saved to {result['output_dir']}") + + def main(): """Main entry point for the CLI.""" parser = argparse.ArgumentParser( @@ -465,6 +624,7 @@ def main(): setup_inference_parser(subparsers) setup_download_parser(subparsers) setup_download_data_parser(subparsers) + setup_train_parser(subparsers) # Parse arguments args = parser.parse_args() @@ -480,6 +640,8 @@ def main(): run_download_cli(args) elif args.command == "download-data": run_download_data_cli(args) + elif args.command == "train": + run_train_cli(args) if __name__ == "__main__": diff --git a/src/transcriptformer/cli/conf/inference_config.yaml b/src/transcriptformer/cli/conf/inference_config.yaml index 37583ab..7351eb2 100644 --- a/src/transcriptformer/cli/conf/inference_config.yaml +++ b/src/transcriptformer/cli/conf/inference_config.yaml @@ -25,7 +25,7 @@ model: emb_type: cell # Type of embeddings to extract: "cell" for mean-pooled cell embeddings or "cge" for contextual gene embeddings num_gpus: 1 # Number of GPUs to use for inference (1 = single GPU, -1 = all available GPUs, >1 = specific number) device: auto # Specific device to use: "auto" (best available), "cpu", "cuda", "mps" - use_oom_dataloader: false # Use OOM-safe map-style DataLoader with DistributedSampler + use_oom_dataloader: false # Use OOM-safe map-style DataLoader (Lightning handles distributed sampling in DDP) data_config: _target_: transcriptformer.data.dataclasses.DataConfig diff --git a/src/transcriptformer/cli/conf/train_config.yaml b/src/transcriptformer/cli/conf/train_config.yaml new file mode 100644 index 0000000..83b57cf --- /dev/null +++ b/src/transcriptformer/cli/conf/train_config.yaml @@ -0,0 +1,88 @@ +defaults: + - _self_ # Load the current config file + +train: + # === Checkpointing & Output === + checkpoint_dir: null # Path to source checkpoint directory (required) + resume_artifact_dir: null # Path to previous training output for resuming (optional) + resume_mode: weights # 'weights' (load weights only) or 'lightning' (resume full state) + output_dir: ./train_outputs/tf_assay_expanded # Output directory for new artifacts + + checkpoint: + save_top_k: 1 # Number of best checkpoints to keep (by val/total_loss) + save_last: true # Always save the last checkpoint + + live_metrics: + enabled: true # Export per-epoch summaries and live plots while training + every_n_epochs: 1 # Refresh metrics/plots every N epochs + + # === Data === + train_files: [] # List of training .h5ad files + val_files: [] # List of validation .h5ad files + expanded_assay_vocab: null # Path to expanded assay vocab JSON (optional) + obs_assay_col: assay # Required AnnData.obs column used to create aux tokens for model forward + obs_keys: null # Optional list of extra AnnData.obs fields to carry in batch.obs (metadata only), use `all` to include all obs columns. + + data_config: + _target_: transcriptformer.data.dataclasses.DataConfig + gene_col_name: ensembl_id # Column name in AnnData.var containing gene identifiers + clip_counts: 30.0 # Maximum count value (higher values will be clipped), Null means inherit from checkpoint model config + filter_to_vocabs: true # Whether to filter genes to only those in the vocabulary, Null means inherit from checkpoint model config + filter_outliers: 3.0 # Standard deviation threshold for filtering outlier cells (0.0 = no filtering), Null means inherit from checkpoint model config + normalize_to_scale: 0.0 # Scale factor for count normalization (0 = no normalization) + sort_genes: null # Whether to sort genes by expression level, Null means inherit from checkpoint model config + randomize_genes: true # Whether to randomize gene order, Null means inherit from checkpoint model config + min_expressed_genes: 200 # Minimum number of expressed genes required per cell, Null means inherit from checkpoint model config + use_raw: "auto" # Whether to use .raw.X (True), .X (False), or auto-detect (auto/null), Null means inherit from checkpoint model config + remove_duplicate_genes: false # Remove duplicate genes instead of raising + n_data_workers: 4 # Data workers used by data module config + + # === Dataloader === + batch_size: 8 # Batch size + num_workers: 4 # DataLoader workers + use_oom_dataloader: true # Use backed OOM dataset + file-aware batching (lower memory, mostly file-local batches) + enable_file_aware_batching: true # Training-only and DDP-safe: keeps file-local OOM batches; set false to use Lightning DistributedSampler path + oom_batches_per_file: 1 # In OOM mode, number of consecutive batches drawn from one file before interleaving + + + # === Training Loop === + # Recommended: Aim for a global batch size of 4–5 million tokens (examples * seq_len * num_gpus * num_nodes * accumulate_grad_batches) + max_epochs: 5 # Number of epochs + precision: 16-mixed # Precision (16-mixed, 32, bf16, etc.) + device: auto # Preferred device policy ('auto', 'cpu', 'cuda', 'mps'), same resolution path as inference + num_gpus: 1 # Requested GPU count (1, >1, or -1 for all available on each node) + num_nodes: 1 # Number of nodes/machines (for DDP) + log_every_n_steps: 10 # Logging frequency + seed: 42 # Random seed + gradient_clip_val: 1.0 # Max norm for gradient clipping (improves stability) + accumulate_grad_batches: 6 # Number of steps to accumulate gradients before optimizer step + + # === Optimizer & Scheduler === + lr: 5.5e-5 # Learning rate, try 1.5e-4 if train from scratch + weight_decay: 0.05 # AdamW weight decay, try 0.1 if train from scratch + adam_beta1: 0.9 # AdamW beta1 + adam_beta2: 0.95 # AdamW beta2 + adam_eps: 1e-8 # AdamW epsilon + warmup_ratio: 0.1 # LR warmup ratio + min_lr_ratio: 0.18 # Minimum LR as fraction of initial, cap at 1e-5 + + # === Loss Weights === + loss_config: + _target_: transcriptformer.data.dataclasses.LossConfig + gene_id_loss_weight: 1.0 # Gene ID prediction loss weight + softplus_approx: true # Use softplus approximation in loss path + + # === Assay Embedding Init === + init_default_source: unknown # Default source for new assay tokens + assay_init_map: [] # List of 'new=source' mappings + + # === Freezing & Training Policy === + freeze_transformer: false # Freeze transformer encoder + unfreeze_last_n_transformer_blocks: 0 # If freeze_transformer=true, unfreeze only the last N encoder blocks (0 keeps full encoder frozen) + freeze_gene_embeddings: true # Freeze gene embeddings, this is pretrained from protein ESM2. + freeze_count_head: false # Freeze count decoder + freeze_gene_head: false # Freeze gene decoder + train_aux_only: false # Train only aux (assay) embeddings + + # === Data Preprocessing === + shuffle_expressed_each_batch: false # Shuffle expressed genes per batch diff --git a/src/transcriptformer/cli/train.py b/src/transcriptformer/cli/train.py new file mode 100644 index 0000000..233b93e --- /dev/null +++ b/src/transcriptformer/cli/train.py @@ -0,0 +1,29 @@ +"""Hydra entrypoint for training Transcriptformer assay adaptations.""" + +from __future__ import annotations + +import logging +import os + +import hydra +from omegaconf import DictConfig, OmegaConf + +from transcriptformer.train.engine import run_train_from_dict, setup_runtime_for_training + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") + + +@hydra.main( + config_path=os.path.join(os.path.dirname(__file__), "conf"), + config_name="train_config.yaml", + version_base=None, +) +def main(cfg: DictConfig): + setup_runtime_for_training() + logging.info("Training config:\n%s", OmegaConf.to_yaml(cfg)) + result = run_train_from_dict(OmegaConf.to_container(cfg.train, resolve=True)) + logging.info("Training complete: %s", result) + + +if __name__ == "__main__": + main() diff --git a/src/transcriptformer/data/dataloader.py b/src/transcriptformer/data/dataloader.py index a436be4..48c437c 100644 --- a/src/transcriptformer/data/dataloader.py +++ b/src/transcriptformer/data/dataloader.py @@ -9,7 +9,7 @@ import torch from scipy.sparse import csc_matrix, csr_matrix from torch import tensor -from torch.utils.data import Dataset +from torch.utils.data import BatchSampler, Dataset from transcriptformer.data.dataclasses import BatchData from transcriptformer.tokenizer.tokenizer import ( @@ -251,6 +251,7 @@ def load_gene_features( gene_counts = Counter(gene_names) duplicates = {gene for gene, count in gene_counts.items() if count > 1} + dedup_col_indices = None if len(duplicates) > 0: if remove_duplicate_genes: seen = set() @@ -259,20 +260,25 @@ def load_gene_features( if gene not in seen: seen.add(gene) unique_indices.append(i) - adata = adata[:, unique_indices].copy() gene_names = gene_names[unique_indices] logging.warning( f"Removed {len(duplicates)} duplicate genes after removing version numbers. Kept first occurrence." ) + if adata.isbacked: + # Cannot copy a backed AnnData object; return indices so the caller + # can incorporate them into its column-index filter instead. + dedup_col_indices = unique_indices + else: + adata = adata[:, unique_indices].copy() else: raise ValueError( "Found duplicate genes after removing version numbers. " "Remove duplicates or pass --remove-duplicate-genes." ) - return gene_names, True, adata + return gene_names, True, adata, dedup_col_indices except KeyError: - return None, False, adata + return None, False, adata, None def validate_gene_dimension(X: np.ndarray, gene_names: np.ndarray, gene_col_name: str): @@ -283,6 +289,43 @@ def validate_gene_dimension(X: np.ndarray, gene_names: np.ndarray, gene_col_name ) +def compute_row_stats_chunked( + X_layer, + filter_idx: list[int] | None = None, + chunk_size: int = 1024, +) -> tuple[np.ndarray, np.ndarray]: + """Compute per-row expression sums and nonzero counts in chunks. + + This avoids loading the full matrix into memory for backed AnnData layers. + """ + n_rows = int(X_layer.shape[0]) + expr_counts = np.zeros(n_rows, dtype=np.float64) + nnz_counts = np.zeros(n_rows, dtype=np.int32) + + for start in range(0, n_rows, chunk_size): + end = min(start + chunk_size, n_rows) + block = X_layer[start:end] + + if isinstance(block, csr_matrix | csc_matrix): + if filter_idx is not None: + block = block[:, filter_idx] + expr = np.asarray(block.sum(axis=1)).ravel() + nnz = np.asarray(block.getnnz(axis=1)).ravel() + else: + arr = np.asarray(block) + if arr.ndim == 1: + arr = arr.reshape(1, -1) + if filter_idx is not None: + arr = arr[:, filter_idx] + expr = arr.sum(axis=1) + nnz = (arr > 0).sum(axis=1) + + expr_counts[start:end] = expr + nnz_counts[start:end] = nnz + + return expr_counts, nnz_counts + + class AnnDataset(Dataset): def __init__( self, @@ -328,6 +371,7 @@ def __init__( self.obs_keys = obs_keys self.use_raw = use_raw self.remove_duplicate_genes = remove_duplicate_genes + self.filter_metadata: list[dict] = [] self.gene_tokenizer = BatchGeneTokenizer(gene_vocab) if aux_vocab is not None: @@ -356,7 +400,7 @@ def _get_batch_from_file(self, file: str | anndata.AnnData) -> BatchData | None: logging.error(f"Failed to load data from {file_path}") return None - gene_names, success, adata = load_gene_features( + gene_names, success, adata, _dedup_col_indices = load_gene_features( adata, self.gene_col_name, self.remove_duplicate_genes, use_raw=self.use_raw ) if not success: @@ -379,6 +423,9 @@ def _get_batch_from_file(self, file: str | anndata.AnnData) -> BatchData | None: "If your data is normalized, consider using the original count matrix instead." ) + original_gene_count = int(len(gene_names)) + original_cell_count = int(X.shape[0]) + logging.info("Applying filters") vocab = self.gene_vocab X, obs, gene_names = apply_filters( @@ -393,9 +440,38 @@ def _get_batch_from_file(self, file: str | anndata.AnnData) -> BatchData | None: ) if X is None: + self.filter_metadata.append( + { + "file": file_path, + "original_genes": original_gene_count, + "kept_genes": 0, + "removed_genes": original_gene_count, + "original_cells": original_cell_count, + "kept_cells": 0, + "removed_cells": original_cell_count, + "filter_to_vocab": bool(self.filter_to_vocab), + "filter_outliers": float(self.filter_outliers), + "min_expressed_genes": int(self.min_expressed_genes), + } + ) logging.warning(f"Data was filtered out completely for {file_path}") return None + self.filter_metadata.append( + { + "file": file_path, + "original_genes": original_gene_count, + "kept_genes": int(len(gene_names)), + "removed_genes": int(original_gene_count - len(gene_names)), + "original_cells": original_cell_count, + "kept_cells": int(X.shape[0]), + "removed_cells": int(original_cell_count - X.shape[0]), + "filter_to_vocab": bool(self.filter_to_vocab), + "filter_outliers": float(self.filter_outliers), + "min_expressed_genes": int(self.min_expressed_genes), + } + ) + logging.info("Processing data") batch = process_batch( X, @@ -525,10 +601,13 @@ def __init__( pad_token: str = "[PAD]", gene_col_name: str = "ensembl_id", filter_to_vocab: bool = True, + filter_outliers: float = 0.0, + min_expressed_genes: int = 0, clip_counts: float = 1e10, obs_keys: list[str] | None = None, use_raw: bool | None = None, remove_duplicate_genes: bool = False, + stats_chunk_size: int = 1024, ): super().__init__() self.files_list = files_list @@ -543,10 +622,14 @@ def __init__( self.pad_token = pad_token self.gene_col_name = gene_col_name self.filter_to_vocab = filter_to_vocab + self.filter_outliers = filter_outliers + self.min_expressed_genes = min_expressed_genes self.clip_counts = clip_counts self.obs_keys = obs_keys self.use_raw = use_raw self.remove_duplicate_genes = remove_duplicate_genes + self.stats_chunk_size = max(1, int(stats_chunk_size)) + self.filter_metadata: list[dict] = [] self.gene_tokenizer = BatchGeneTokenizer(gene_vocab) if aux_vocab is not None: @@ -556,69 +639,155 @@ def __init__( self._handles: list[anndata.AnnData] = [] self._gene_names_per_file: list[np.ndarray] = [] self._filter_idx_per_file: list[list[int] | None] = [] + self._keep_rows_per_file: list[np.ndarray] = [] self._X_per_file: list = [] self._n_rows: list[int] = [] for file in self.files_list: file_path = file if self.data_dir is None else os.path.join(self.data_dir, file) adata = anndata.read_h5ad(file_path, backed="r") - gene_names, success, adata = load_gene_features( + gene_names, success, adata, dedup_col_indices = load_gene_features( adata, self.gene_col_name, self.remove_duplicate_genes, use_raw=self.use_raw ) if not success: raise ValueError(f"Failed to load gene features from {file_path}") - # Optional vocab filtering at token level + + original_gene_count = int(len(gene_names)) + original_cell_count = int(adata.n_obs) + + # Optional vocab filtering at token level. + # When dedup_col_indices is set (backed mode deduplication), those indices are + # into the *original* adata column space; vocab filtering indices are into the + # deduplicated gene_names list, so the two must be composed. filter_idx = None if self.filter_to_vocab: - original_gene_count = len(gene_names) - filter_idx = [i for i, name in enumerate(gene_names) if name in self.gene_vocab] - gene_names = gene_names[filter_idx] + vocab_positions = [i for i, name in enumerate(gene_names) if name in self.gene_vocab] + if dedup_col_indices is not None: + filter_idx = [dedup_col_indices[i] for i in vocab_positions] + else: + filter_idx = vocab_positions + gene_names = gene_names[np.array(vocab_positions)] logging.info( f"Filtered {original_gene_count} genes to {len(gene_names)} genes in vocab for file {file_path}" ) if len(gene_names) == 0: raise ValueError(f"No genes remaining after filtering for file {file_path}") + elif dedup_col_indices is not None: + # No vocab filter but deduplication was deferred; use the dedup indices as filter_idx + filter_idx = dedup_col_indices + + X_layer = get_counts_layer(adata, self.use_raw) + if self.filter_outliers > 0 or self.min_expressed_genes > 0: + expr_counts, nnz_counts = compute_row_stats_chunked( + X_layer, + filter_idx=filter_idx, + chunk_size=self.stats_chunk_size, + ) + + keep_mask = np.ones(original_cell_count, dtype=bool) + if self.filter_outliers > 0: + count_std = np.std(expr_counts) + count_mean = np.mean(expr_counts) + keep_mask &= (expr_counts > count_mean - count_std * self.filter_outliers) & ( + expr_counts < count_mean + count_std * self.filter_outliers + ) + if self.min_expressed_genes > 0: + keep_mask &= nnz_counts >= self.min_expressed_genes + + keep_rows = np.nonzero(keep_mask)[0].astype(np.int64) + + logging.info( + f"Filtered {original_cell_count} cells to {len(keep_rows)} cells for file {file_path}" + ) + else: + keep_rows = np.arange(original_cell_count, dtype=np.int64) + if keep_rows.size == 0: + raise ValueError(f"No cells remaining after filtering for file {file_path}") self._handles.append(adata) self._gene_names_per_file.append(gene_names) self._filter_idx_per_file.append(filter_idx) - X_layer = get_counts_layer(adata, self.use_raw) + self._keep_rows_per_file.append(keep_rows) self._X_per_file.append(X_layer) - self._n_rows.append(int(adata.n_obs)) + self._n_rows.append(int(keep_rows.size)) + + self.filter_metadata.append( + { + "file": file_path, + "original_genes": original_gene_count, + "kept_genes": int(len(gene_names)), + "removed_genes": int(original_gene_count - len(gene_names)), + "original_cells": original_cell_count, + "kept_cells": int(keep_rows.size), + "removed_cells": int(original_cell_count - keep_rows.size), + "filter_to_vocab": bool(self.filter_to_vocab), + "filter_outliers": float(self.filter_outliers), + "min_expressed_genes": int(self.min_expressed_genes), + } + ) self._offsets = np.cumsum([0] + self._n_rows) def __len__(self) -> int: return int(self._offsets[-1]) + @property + def num_files(self) -> int: + return len(self._n_rows) + + def file_num_rows(self, file_id: int) -> int: + return int(self._n_rows[file_id]) + + def file_offset(self, file_id: int) -> int: + return int(self._offsets[file_id]) + def _loc(self, idx: int) -> tuple[int, int]: file_id = int(np.searchsorted(self._offsets, idx, side="right") - 1) row = int(idx - self._offsets[file_id]) return file_id, row - def __getitem__(self, idx: int) -> BatchData: - file_id, row = self._loc(idx) - adata = self._handles[file_id] - gene_names = self._gene_names_per_file[file_id] - filter_idx = self._filter_idx_per_file[file_id] + @staticmethod + def _ensure_2d_array(x_rows) -> np.ndarray: + if isinstance(x_rows, csr_matrix | csc_matrix): + arr = x_rows.toarray() + else: + arr = np.asarray(x_rows) + if arr.ndim == 1: + arr = arr.reshape(1, -1) + return arr + + def _build_obs_dict(self, obs_rows) -> dict[str, np.ndarray] | None: + if self.obs_keys is None: + return None + + obs_dict = {} + cols = list(obs_rows.columns) if "all" in self.obs_keys else list(self.obs_keys or []) + for col in cols: + obs_dict[col] = np.array(obs_rows[col].tolist())[:, None] + return obs_dict + + def _build_batch_from_rows(self, file_id: int, rows: list[int]) -> BatchData: + actual_rows = self._keep_rows_per_file[file_id][np.asarray(rows, dtype=np.int64)] + actual_rows = np.asarray(actual_rows, dtype=np.int64) X = self._X_per_file[file_id] - # Some backed sparse implementations use __getitem__ returning 2D; ensure 1D - x_row = X[row] - # Only convert to dense if the row is actually sparse - if isinstance(x_row, csr_matrix | csc_matrix): - x_row = x_row.toarray().ravel() - else: - x_row = np.asarray(x_row).ravel() + # Backed/h5ad row reads are typically better behaved (and often faster) with monotonic indices + sort_order = np.argsort(actual_rows, kind="stable") + sorted_rows = actual_rows[sort_order] + x_rows = self._ensure_2d_array(X[sorted_rows]) + if x_rows.shape[0] > 1: + x_rows = x_rows[np.argsort(sort_order)] + + filter_idx = self._filter_idx_per_file[file_id] if filter_idx is not None: - x_row = x_row[filter_idx] + x_rows = x_rows[:, filter_idx] - obs_row = adata.obs.iloc[row : row + 1] + adata = self._handles[file_id] + obs_rows = adata.obs.iloc[actual_rows] + gene_names = self._gene_names_per_file[file_id] - # Build a 1-row batch and reuse existing processing pipeline - x_batch = np.expand_dims(x_row, axis=0) batch = process_batch( - x_batch, - obs_row, + x_rows, + obs_rows, gene_names, self.gene_tokenizer, getattr(self, "aux_tokenizer", None), @@ -633,20 +802,151 @@ def __getitem__(self, idx: int) -> BatchData: self.aux_vocab, ) - # Convert to BatchData for collate_fn compatibility - obs_dict = None - if self.obs_keys is not None: - obs_dict = {} - cols = list(obs_row.columns) if "all" in self.obs_keys else list(self.obs_keys or []) - for col in cols: - obs_dict[col] = np.array(obs_row[col].tolist())[:, None] + return BatchData( + gene_counts=batch["gene_counts"], + gene_token_indices=batch["gene_token_indices"], + file_path=None, + aux_token_indices=batch.get("aux_token_indices"), + obs=self._build_obs_dict(obs_rows), + ) + def __getitem__(self, idx: int) -> BatchData: + # Build a 1-row batch + file_id, row = self._loc(idx) + batch = self._build_batch_from_rows(file_id, [row]) return BatchData( - gene_counts=batch["gene_counts"][0], - gene_token_indices=batch["gene_token_indices"][0], + gene_counts=batch.gene_counts[0], + gene_token_indices=batch.gene_token_indices[0], file_path=None, - aux_token_indices=( - batch.get("aux_token_indices")[0] if batch.get("aux_token_indices") is not None else None - ), - obs=obs_dict, + aux_token_indices=(batch.aux_token_indices[0] if batch.aux_token_indices is not None else None), + obs=({col: value[0:1] for col, value in batch.obs.items()} if batch.obs is not None else None), ) + + def __getitems__(self, indices: list[int]) -> BatchData | list[BatchData]: + # Build a multi-row batch if all indices are from the same file; otherwise fallback to individual retrieval + if not indices: + return [] + + locations = [self._loc(int(idx)) for idx in indices] + file_ids = {file_id for file_id, _ in locations} + if len(file_ids) != 1: + return [self[int(idx)] for idx in indices] + + file_id = locations[0][0] + rows = [row for _, row in locations] + return self._build_batch_from_rows(file_id, rows) + + +class FileAwareBatchSampler(BatchSampler): + """Yield file-local batches while interleaving files across an epoch.""" + + def __init__( + self, + dataset: AnnDatasetOOM, + batch_size: int, + shuffle: bool = True, + drop_last: bool = False, + batches_per_file: int = 1, + seed: int = 0, + ): + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if batches_per_file <= 0: + raise ValueError("batches_per_file must be positive") + + self.dataset = dataset + self.batch_size = int(batch_size) + self.shuffle = bool(shuffle) + self.drop_last = bool(drop_last) + self.batches_per_file = int(batches_per_file) + self.seed = int(seed) + self._epoch = 0 + + def set_epoch(self, epoch: int) -> None: + self._epoch = int(epoch) + + def _build_all_batches(self) -> list[list[int]]: + rng = random.Random(self.seed + self._epoch) + file_ids = list(range(self.dataset.num_files)) + if self.shuffle: + rng.shuffle(file_ids) + + # Validation-friendly path: deterministic, file-sequential traversal. + # This minimizes backed-file switching and is typically faster than interleaving. + if not self.shuffle: + all_batches: list[list[int]] = [] + for file_id in file_ids: + local_rows = list(range(self.dataset.file_num_rows(file_id))) + offset = self.dataset.file_offset(file_id) + for start in range(0, len(local_rows), self.batch_size): + batch_rows = local_rows[start : start + self.batch_size] + if len(batch_rows) < self.batch_size and self.drop_last: + continue + all_batches.append([offset + row for row in batch_rows]) + return all_batches + + active_files: list[dict] = [] + for file_id in file_ids: + local_rows = list(range(self.dataset.file_num_rows(file_id))) + if self.shuffle: + rng.shuffle(local_rows) + + batches = [] + offset = self.dataset.file_offset(file_id) + for start in range(0, len(local_rows), self.batch_size): + batch_rows = local_rows[start : start + self.batch_size] + if len(batch_rows) < self.batch_size and self.drop_last: + continue + batches.append([offset + row for row in batch_rows]) + + if batches: + active_files.append({"batches": batches, "cursor": 0}) + + all_batches: list[list[int]] = [] + while active_files: + round_order = list(range(len(active_files))) + if self.shuffle: + rng.shuffle(round_order) + + next_active = [] + for pos in round_order: + state = active_files[pos] + start = state["cursor"] + stop = min(start + self.batches_per_file, len(state["batches"])) + all_batches.extend(state["batches"][start:stop]) + if stop < len(state["batches"]): + state["cursor"] = stop + next_active.append(state) + active_files = next_active + + return all_batches + + def _rank_batches(self, all_batches: list[list[int]]) -> list[list[int]]: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return all_batches + + world_size = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() + if world_size <= 1 or not all_batches: + return all_batches + + if self.drop_last: + total = len(all_batches) - (len(all_batches) % world_size) + all_batches = all_batches[:total] + else: + remainder = len(all_batches) % world_size + if remainder: + padding = world_size - remainder + all_batches = all_batches + all_batches[:padding] + + return all_batches[rank::world_size] + + def __iter__(self): + rank_batches = self._rank_batches(self._build_all_batches()) + self._epoch += 1 + yield from rank_batches + + def __len__(self) -> int: + all_batches = self._build_all_batches() + rank_batches = self._rank_batches(all_batches) + return len(rank_batches) diff --git a/src/transcriptformer/model/assay_adaptation.py b/src/transcriptformer/model/assay_adaptation.py new file mode 100644 index 0000000..0453d43 --- /dev/null +++ b/src/transcriptformer/model/assay_adaptation.py @@ -0,0 +1,139 @@ +"""Utilities for adapting Transcriptformer to expanded assay vocabularies.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn + + +@dataclass +class AssayInitConfig: + """Configuration for initializing new assay embedding rows.""" + + default_source: str = "unknown" + mean_pool_fallback: bool = True + + +def _validate_vocab_indices(vocab: dict[str, int]) -> None: + ids = set(vocab.values()) + expected = set(range(len(vocab))) + if ids != expected: + raise ValueError("Assay vocab indices must be contiguous [0, ..., N-1]") + + +def build_expanded_assay_embedding_weight( + old_weight: torch.Tensor, + old_vocab: dict[str, int], + new_vocab: dict[str, int], + init_map: dict[str, str] | None = None, + cfg: AssayInitConfig | None = None, +) -> torch.Tensor: + """Create an expanded embedding matrix for assay tokens. + + Args: + old_weight: Existing assay embedding matrix with shape [old_vocab_size, dim]. + old_vocab: Original assay vocab mapping token -> index. + new_vocab: Expanded assay vocab mapping token -> index. + init_map: Optional explicit mapping new_token -> existing_source_token. + cfg: Initialization policy. + + Returns + ------- + Expanded embedding matrix with shape [new_vocab_size, dim]. + """ + _validate_vocab_indices(old_vocab) + _validate_vocab_indices(new_vocab) + + # Fast path: unchanged vocab mapping means weights can be reused as-is. + if old_vocab == new_vocab: + return old_weight.clone() + + if cfg is None: + cfg = AssayInitConfig() + if init_map is None: + init_map = {} + + dim = old_weight.shape[1] + device = old_weight.device + dtype = old_weight.dtype + + new_weight = torch.empty((len(new_vocab), dim), device=device, dtype=dtype) + + # Random init baseline similar to nn.Embedding default behavior. + nn.init.normal_(new_weight, mean=0.0, std=1.0) + + # Precompute fallback vector. + if cfg.mean_pool_fallback: + fallback_vec = old_weight.mean(dim=0) + else: + source_token = cfg.default_source + if source_token not in old_vocab: + raise ValueError(f"default_source token '{source_token}' is missing in old vocab") + fallback_vec = old_weight[old_vocab[source_token]] + + for token, new_idx in new_vocab.items(): + if token in old_vocab: + new_weight[new_idx] = old_weight[old_vocab[token]] + continue + + source_token = init_map.get(token, cfg.default_source) + if source_token in old_vocab: + new_weight[new_idx] = old_weight[old_vocab[source_token]] + else: + new_weight[new_idx] = fallback_vec + + return new_weight + + +def apply_freeze_policy( + model: nn.Module, + freeze_transformer: bool = False, + freeze_gene_embeddings: bool = True, + freeze_count_head: bool = False, + freeze_gene_head: bool = False, + train_aux_only: bool = False, + unfreeze_last_n_transformer_blocks: int = 0, +) -> None: + """Apply parameter freezing policy in-place.""" + + if train_aux_only: + for param in model.parameters(): + param.requires_grad = False + if hasattr(model, "aux_embeddings"): + for param in model.aux_embeddings.parameters(): + param.requires_grad = True + return + + if freeze_transformer and hasattr(model, "transformer_encoder"): + for param in model.transformer_encoder.parameters(): + param.requires_grad = False + + # Optional adapter-style policy: keep the encoder mostly frozen, + # but allow updates in the last N transformer blocks. + n = max(0, int(unfreeze_last_n_transformer_blocks)) + if n > 0 and hasattr(model.transformer_encoder, "encoder_layers"): + layers = model.transformer_encoder.encoder_layers + for layer in list(layers)[-n:]: + for param in layer.parameters(): + param.requires_grad = True + + if freeze_gene_embeddings and hasattr(model, "gene_embeddings"): + for param in model.gene_embeddings.parameters(): + param.requires_grad = False + + if freeze_count_head and hasattr(model, "mu"): + for param in model.mu.parameters(): + param.requires_grad = False + + if freeze_gene_head and hasattr(model, "gene_id_head"): + for param in model.gene_id_head.parameters(): + param.requires_grad = False + + +def count_trainable_parameters(model: nn.Module) -> tuple[int, int]: + """Return (trainable_params, total_params).""" + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return trainable, total diff --git a/src/transcriptformer/model/inference.py b/src/transcriptformer/model/inference.py index f61e8a8..08c5333 100644 --- a/src/transcriptformer/model/inference.py +++ b/src/transcriptformer/model/inference.py @@ -12,6 +12,7 @@ from transcriptformer.data.dataloader import AnnDataset, AnnDatasetOOM from transcriptformer.model.embedding_surgery import change_embedding_layer from transcriptformer.tokenizer.vocab import load_vocabs_and_embeddings +from transcriptformer.utils.device import resolve_accelerator_and_devices, resolve_checkpoint_map_location from transcriptformer.utils.utils import stack_dict # Set float32 matmul precision for better performance with Tensor Cores @@ -83,21 +84,8 @@ def run_inference(cfg, data_files: list[str] | list[anndata.AnnData]): ) logging.info("Loading model checkpoint") - # Determine target device for loading checkpoint device_preference = getattr(cfg.model.inference_config, "device", "auto") - if device_preference == "cpu": - map_location = "cpu" - elif device_preference == "cuda": - map_location = "cuda" if torch.cuda.is_available() else "cpu" - elif device_preference == "mps": - map_location = "mps" if torch.backends.mps.is_available() else "cpu" - else: # auto - if torch.cuda.is_available(): - map_location = "cuda" - elif torch.backends.mps.is_available(): - map_location = "mps" - else: - map_location = "cpu" + map_location = resolve_checkpoint_map_location(device_preference) # Instead of loading full checkpoint, just load weights with proper device mapping state_dict = torch.load(cfg.model.inference_config.load_checkpoint, weights_only=True, map_location=map_location) @@ -136,8 +124,9 @@ def run_inference(cfg, data_files: list[str] | list[anndata.AnnData]): "remove_duplicate_genes": cfg.model.data_config.remove_duplicate_genes, "use_raw": cfg.model.data_config.use_raw, } - if getattr(cfg.model.inference_config, "use_oom_dataloader", False): - # Use OOM-safe map-style dataset + use_oom = bool(getattr(cfg.model.inference_config, "use_oom_dataloader", False)) + can_use_oom = all(isinstance(file, str) for file in data_files) + if use_oom and can_use_oom: dataset = AnnDatasetOOM( data_files, gene_vocab, @@ -155,6 +144,10 @@ def run_inference(cfg, data_files: list[str] | list[anndata.AnnData]): remove_duplicate_genes=cfg.model.data_config.remove_duplicate_genes, ) else: + if use_oom and not can_use_oom: + logging.warning( + "use_oom_dataloader=true with in-memory AnnData inputs; falling back to AnnDataset for compatibility" + ) dataset = AnnDataset(data_files, **data_kwargs) # Create dataloader @@ -167,71 +160,8 @@ def run_inference(cfg, data_files: list[str] | list[anndata.AnnData]): collate_fn=dataset.collate_fn, ) - # Determine device and accelerator based on user preference - device_preference = getattr(cfg.model.inference_config, "device", "auto") num_gpus = getattr(cfg.model.inference_config, "num_gpus", 1) - - # Handle device preference - if device_preference == "cpu": - # Force CPU usage - accelerator = "cpu" - devices = 1 - logging.info("Forcing CPU usage based on device preference") - elif device_preference == "cuda": - # Force CUDA usage - if not torch.cuda.is_available(): - raise RuntimeError("CUDA requested but not available") - accelerator = "gpu" - if num_gpus == -1: - devices = torch.cuda.device_count() - elif num_gpus > 1: - available_gpus = torch.cuda.device_count() - if available_gpus < num_gpus: - logging.warning( - f"Requested {num_gpus} CUDA devices but only {available_gpus} available. Using {available_gpus}." - ) - devices = available_gpus - else: - devices = num_gpus - else: - devices = 1 - logging.info(f"Forcing CUDA usage with {devices} device(s)") - elif device_preference == "mps": - # Force MPS usage (Apple Silicon) - if not torch.backends.mps.is_available(): - raise RuntimeError("MPS requested but not available") - accelerator = "mps" - devices = 1 # MPS typically supports single device - logging.info("Forcing MPS usage for Apple Silicon") - else: # device_preference == "auto" - # Auto-detect best available device (existing logic) - if num_gpus == -1: - # Use all available GPUs - devices = torch.cuda.device_count() if torch.cuda.is_available() else 1 - accelerator = ( - "gpu" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") - ) - elif num_gpus > 1: - # Use specified number of GPUs - available_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0 - if available_gpus < num_gpus: - logging.warning( - f"Requested {num_gpus} GPUs but only {available_gpus} available. Using {available_gpus} GPUs." - ) - devices = available_gpus if available_gpus > 0 else 1 - accelerator = "gpu" if available_gpus > 0 else ("mps" if torch.backends.mps.is_available() else "cpu") - else: - devices = num_gpus - accelerator = "gpu" - else: - # Use single GPU or CPU - devices = 1 - if torch.cuda.is_available(): - accelerator = "gpu" - elif torch.backends.mps.is_available(): - accelerator = "mps" - else: - accelerator = "cpu" + accelerator, devices = resolve_accelerator_and_devices(device_preference, num_gpus) logging.info(f"Using {devices} device(s) with accelerator: {accelerator}") diff --git a/src/transcriptformer/train/__init__.py b/src/transcriptformer/train/__init__.py new file mode 100644 index 0000000..dc6b70d --- /dev/null +++ b/src/transcriptformer/train/__init__.py @@ -0,0 +1,12 @@ +"""Training utilities for Transcriptformer.""" + +from transcriptformer.train.assay_vocab import discover_tokens, expand_vocab, load_vocab, write_vocab +from transcriptformer.train.engine import run_train_from_dict + +__all__ = [ + "discover_tokens", + "expand_vocab", + "load_vocab", + "write_vocab", + "run_train_from_dict", +] diff --git a/src/transcriptformer/train/assay_vocab.py b/src/transcriptformer/train/assay_vocab.py new file mode 100644 index 0000000..ff8ed23 --- /dev/null +++ b/src/transcriptformer/train/assay_vocab.py @@ -0,0 +1,48 @@ +"""Assay vocabulary helpers for adaptation workflows.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import anndata as ad + + +def load_vocab(path: Path) -> dict[str, int]: + with path.open() as f: + vocab = json.load(f) + if "unknown" not in vocab: + raise ValueError("Input vocab must include an 'unknown' token") + return vocab + + +def write_vocab(vocab: dict[str, int], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w") as f: + json.dump(vocab, f, indent=2, sort_keys=True) + f.write("\n") + + +def discover_tokens(data_files: list[str], obs_col: str, min_frequency: int = 1) -> list[str]: + counts: dict[str, int] = {} + for file in data_files: + adata = ad.read_h5ad(file, backed="r") + if obs_col not in adata.obs.columns: + raise ValueError(f"Column '{obs_col}' not found in {file}") + values = adata.obs[obs_col].astype(str).values + for token in values: + counts[token] = counts.get(token, 0) + 1 + return sorted([token for token, n in counts.items() if n >= min_frequency]) + + +def expand_vocab(existing_vocab: dict[str, int], new_tokens: list[str]) -> dict[str, int]: + expanded = dict(existing_vocab) + next_id = max(expanded.values()) + 1 if expanded else 0 + + for token in new_tokens: + if token in expanded: + continue + expanded[token] = next_id + next_id += 1 + + return expanded diff --git a/src/transcriptformer/train/engine.py b/src/transcriptformer/train/engine.py new file mode 100644 index 0000000..546a6a9 --- /dev/null +++ b/src/transcriptformer/train/engine.py @@ -0,0 +1,778 @@ +"""Training engine for assay-vocab expansion fine-tuning.""" + +from __future__ import annotations + +import csv +import json +import math +import os +import shutil +from pathlib import Path + +import pytorch_lightning as pl +import torch +from pytorch_lightning.callbacks import ModelCheckpoint +from torch.optim import AdamW +from torch.optim.lr_scheduler import LambdaLR +from torch.utils.data import DataLoader + +from transcriptformer.data.dataclasses import BatchData, DataConfig, LossConfig, ModelConfig +from transcriptformer.data.dataloader import AnnDataset, AnnDatasetOOM, FileAwareBatchSampler +from transcriptformer.model.assay_adaptation import ( + AssayInitConfig, + apply_freeze_policy, + build_expanded_assay_embedding_weight, + count_trainable_parameters, +) +from transcriptformer.model.model import Transcriptformer +from transcriptformer.tokenizer.vocab import construct_gene_embeddings, open_vocabs +from transcriptformer.utils.device import resolve_accelerator_and_devices + + +class LiveMetricsSummaryCallback(pl.Callback): + """Export per-epoch metric summaries and live plots during training.""" + + def __init__(self, output_dir: Path, enabled: bool = True, every_n_epochs: int = 1): + super().__init__() + self.output_dir = Path(output_dir) + self.enabled = enabled + self.every_n_epochs = max(1, int(every_n_epochs)) + self.live_dir = self.output_dir / "live_metrics" + self.epoch_csv_path = self.live_dir / "epoch_metrics.csv" + + @staticmethod + def _to_float(value): + if value is None: + return None + if isinstance(value, torch.Tensor): + value = value.detach().cpu().item() + try: + return float(value) + except (TypeError, ValueError): + return None + + @staticmethod + def _safe_float(text: str | None): + if text is None or text == "": + return None + try: + return float(text) + except ValueError: + return None + + @staticmethod + def _extract_csv_logger_path(trainer: pl.Trainer) -> Path | None: + logger = trainer.logger + if logger is None: + return None + if hasattr(logger, "log_dir"): + return Path(logger.log_dir) / "metrics.csv" + if hasattr(logger, "_logger_iterable"): + for lg in logger._logger_iterable: # Lightning internal logger collection + if hasattr(lg, "log_dir"): + return Path(lg.log_dir) / "metrics.csv" + return None + + def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str) -> None: + if not self.enabled: + return + self.live_dir.mkdir(parents=True, exist_ok=True) + + def on_train_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: + # For training without validation, emit at train epoch end. + if trainer.val_dataloaders: + return + self._export(trainer) + + def on_validation_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: + # For training with validation, emit once per epoch after validation. + self._export(trainer) + + def _export(self, trainer: pl.Trainer) -> None: + if not self.enabled: + return + epoch_idx = int(trainer.current_epoch) + if (epoch_idx + 1) % self.every_n_epochs != 0: + return + + callback_metrics = trainer.callback_metrics + row = { + "epoch": epoch_idx, + "global_step": int(trainer.global_step), + "train/total_loss": self._to_float(callback_metrics.get("train/total_loss")), + "val/total_loss": self._to_float(callback_metrics.get("val/total_loss")), + "train/count_loss": self._to_float(callback_metrics.get("train/count_loss")), + "val/count_loss": self._to_float(callback_metrics.get("val/count_loss")), + "train/gene_loss": self._to_float(callback_metrics.get("train/gene_loss")), + "val/gene_loss": self._to_float(callback_metrics.get("val/gene_loss")), + } + + file_exists = self.epoch_csv_path.exists() + with self.epoch_csv_path.open("a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=list(row.keys())) + if not file_exists: + writer.writeheader() + writer.writerow(row) + + self._render_epoch_bar_plot() + self._render_step_curve_plot(trainer) + + def _render_epoch_bar_plot(self) -> None: + try: + import matplotlib.pyplot as plt + except Exception: + return + + if not self.epoch_csv_path.exists(): + return + + with self.epoch_csv_path.open("r", newline="") as f: + rows = list(csv.DictReader(f)) + if not rows: + return + + last = rows[-1] + labels = [] + values = [] + for key in ["train/total_loss", "val/total_loss", "train/count_loss", "val/count_loss", "train/gene_loss", "val/gene_loss"]: + value = self._safe_float(last.get(key)) + if value is None: + continue + labels.append(key) + values.append(value) + + if not labels: + return + + fig, ax = plt.subplots(figsize=(10, 4)) + ax.bar(labels, values) + ax.set_ylabel("Loss") + ax.set_title(f"Epoch {last.get('epoch', '?')} Metrics") + ax.tick_params(axis="x", rotation=35) + fig.tight_layout() + fig.savefig(self.live_dir / "epoch_metrics_bar_latest.png", dpi=140) + plt.close(fig) + + def _render_step_curve_plot(self, trainer: pl.Trainer) -> None: + csv_metrics = self._extract_csv_logger_path(trainer) + if csv_metrics is None or not csv_metrics.exists(): + return + + try: + import matplotlib.pyplot as plt + except Exception: + return + + with csv_metrics.open("r", newline="") as f: + rows = list(csv.DictReader(f)) + if not rows: + return + + step_key = "step" + candidate_metric_keys = [ + "train/total_loss_step", + "train/count_loss_step", + "train/gene_loss_step", + ] + + # Keep only metrics that actually exist in this run's CSV. + metric_keys = [k for k in candidate_metric_keys if any((r.get(k) not in (None, "")) for r in rows)] + if not metric_keys: + return + + curves = {k: {"x": [], "y": []} for k in metric_keys} + for r in rows: + step_val = self._safe_float(r.get(step_key)) + if step_val is None: + continue + for k in metric_keys: + y = self._safe_float(r.get(k)) + if y is None: + continue + curves[k]["x"].append(step_val) + curves[k]["y"].append(y) + + fig, ax = plt.subplots(figsize=(10, 4)) + have_curve = False + for k in metric_keys: + xs = curves[k]["x"] + ys = curves[k]["y"] + if not xs: + continue + ax.plot(xs, ys, label=k) + have_curve = True + + if not have_curve: + plt.close(fig) + return + + ax.set_xlabel("Step") + ax.set_ylabel("Loss") + ax.set_title("Training Step Metrics") + ax.legend(loc="best") + fig.tight_layout() + fig.savefig(self.live_dir / "step_metrics_curve.png", dpi=140) + plt.close(fig) + + +def _parse_init_map(items: list[str] | None) -> dict[str, str]: + mapping: dict[str, str] = {} + for item in items or []: + if "=" not in item: + raise ValueError(f"Invalid assay init map '{item}', expected new=source") + new_tok, src_tok = item.split("=", 1) + new_tok = new_tok.strip() + src_tok = src_tok.strip() + if not new_tok or not src_tok: + raise ValueError(f"Invalid assay init map '{item}'") + mapping[new_tok] = src_tok + return mapping + + +def _sanitize_data_config(base_data_cfg: dict, cfg: dict) -> DataConfig: + if "data_config" not in cfg or not isinstance(cfg["data_config"], dict): + raise ValueError("Missing required nested 'data_config' in training config") + + out = dict(base_data_cfg) + data_cfg = dict(cfg["data_config"]) + data_cfg.pop("_target_", None) + # Only override base checkpoint data config when values are explicitly provided. + for key, value in data_cfg.items(): + if value is not None: + out[key] = value + + out["aux_vocab_path"] = cfg["source_vocab_dir"] + out["pin_memory"] = True + out["aux_cols"] = cfg.get("obs_assay_col", "assay") + out["pad_zeros"] = True + out["n_data_workers"] = int(out.get("n_data_workers", cfg.get("num_workers", 4))) + out["gene_pad_token"] = out.get("gene_pad_token") or "[PAD]" + out["aux_pad_token"] = out.get("aux_pad_token") or "unknown" + + out.pop("_target_", None) + return DataConfig(**out) + + +def _sanitize_model_config(base_model_cfg: dict) -> ModelConfig: + cfg = dict(base_model_cfg) + cfg.pop("_target_", None) + return ModelConfig(**cfg) + + +def _sanitize_loss_config(base_loss_cfg: dict, cfg: dict) -> LossConfig: + out = dict(base_loss_cfg) + train_loss_cfg = cfg.get("loss_config", {}) + if not isinstance(train_loss_cfg, dict): + raise ValueError("'loss_config' must be a nested mapping in training config") + train_loss_cfg = dict(train_loss_cfg) + train_loss_cfg.pop("_target_", None) + for key, value in train_loss_cfg.items(): + if value is not None: + out[key] = value + out.pop("_target_", None) + return LossConfig(**out) + + +def _shuffle_expressed_genes(batch: BatchData, pad_idx: int) -> BatchData: + tokens = batch.gene_token_indices + counts = batch.gene_counts + + shuffled_tokens = tokens.clone() + shuffled_counts = counts.clone() + + valid = tokens != pad_idx + for i in range(tokens.shape[0]): + idx = torch.nonzero(valid[i], as_tuple=True)[0] + if idx.numel() <= 1: + continue + perm = idx[torch.randperm(idx.numel(), device=idx.device)] + shuffled_tokens[i, idx] = tokens[i, perm] + shuffled_counts[i, idx] = counts[i, perm] + + return BatchData( + gene_counts=shuffled_counts, + gene_token_indices=shuffled_tokens, + aux_token_indices=batch.aux_token_indices, + file_path=batch.file_path, + obs=batch.obs, + ) + + +class TranscriptformerTrainModule(pl.LightningModule): + def __init__( + self, + model: Transcriptformer, + lr: float, + weight_decay: float, + beta1: float, + beta2: float, + eps: float, + warmup_ratio: float, + min_lr_ratio: float, + shuffle_expressed_each_batch: bool, + ): + super().__init__() + self.model = model + self.lr = lr + self.weight_decay = weight_decay + self.beta1 = beta1 + self.beta2 = beta2 + self.eps = eps + self.warmup_ratio = warmup_ratio + self.min_lr_ratio = min_lr_ratio + self.shuffle_expressed_each_batch = shuffle_expressed_each_batch + + def _shared_step(self, batch: BatchData, stage: str) -> torch.Tensor: + if stage == "train" and self.shuffle_expressed_each_batch: + batch = _shuffle_expressed_genes(batch, self.model.gene_vocab.pad_idx) + + out = self.model(batch) + + count_loss = self.model.criterion(mu=out["mu"], input_counts=out["input_counts"], mask=out["mask"]) + gene_loss = torch.tensor(0.0, device=count_loss.device) + if self.model.loss_config.gene_id_loss_weight > 0 and "gene_logit" in out: + gene_loss = self.model.gene_id_criterion( + logits=out["gene_logit"], + input_ids=out["input_gene_token_indices"], + mask=out["mask"], + ) + + total = count_loss + self.model.loss_config.gene_id_loss_weight * gene_loss + # explict batch size for logging to avoid warning from nested batch dictionary + bs = int(batch.gene_counts.shape[0]) + self.log(f"{stage}/total_loss", total, prog_bar=True, on_step=(stage == "train"), on_epoch=True, batch_size=bs) + self.log(f"{stage}/count_loss", count_loss, on_step=False, on_epoch=True, batch_size=bs) + self.log(f"{stage}/gene_loss", gene_loss, on_step=False, on_epoch=True, batch_size=bs) + return total + + def training_step(self, batch: BatchData, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, stage="train") + + def validation_step(self, batch: BatchData, batch_idx: int) -> torch.Tensor: + return self._shared_step(batch, stage="val") + + def configure_optimizers(self): + params = [p for p in self.model.parameters() if p.requires_grad] + if not params: + raise RuntimeError("No trainable parameters found") + + optimizer = AdamW( + params, + lr=self.lr, + betas=(self.beta1, self.beta2), + eps=self.eps, + weight_decay=self.weight_decay, + ) + + total_steps = max(1, int(self.trainer.estimated_stepping_batches)) + warmup_steps = int(total_steps * self.warmup_ratio) + + def lr_lambda(step: int) -> float: + if step < warmup_steps: + return float(step + 1) / float(max(1, warmup_steps)) + progress = (step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) + return self.min_lr_ratio + (1.0 - self.min_lr_ratio) * cosine + + scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda) + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1}, + } + + +def _build_dataset( + files: list[str], + cfg: dict, + data_config: DataConfig, + gene_vocab: dict, + aux_vocab: dict, + seq_len: int, + is_train: bool, +): + for file in files: + if not Path(file).exists(): + raise FileNotFoundError(f"Dataset file not found: {file}") + + kwargs = { + "gene_vocab": gene_vocab, + "aux_vocab": aux_vocab, + "max_len": seq_len, + "normalize_to_scale": data_config.normalize_to_scale, + "sort_genes": bool(data_config.sort_genes), + "randomize_order": bool(data_config.randomize_genes), + "pad_zeros": True, + "gene_col_name": data_config.gene_col_name, + "filter_to_vocab": bool(data_config.filter_to_vocabs), + "clip_counts": data_config.clip_counts, + "use_raw": data_config.use_raw, + "remove_duplicate_genes": bool(data_config.remove_duplicate_genes), + # Optional passthrough obs columns retained in BatchData.obs for downstream hooks/debugging. + "obs_keys": cfg.get("obs_keys", None), + } + + if cfg.get("use_oom_dataloader", False): + return AnnDatasetOOM( + files_list=files, + filter_outliers=float(data_config.filter_outliers), + min_expressed_genes=int(data_config.min_expressed_genes), + **kwargs, + ) + + return AnnDataset( + files_list=files, + filter_outliers=float(data_config.filter_outliers), + min_expressed_genes=int(data_config.min_expressed_genes), + seed=int(cfg.get("seed", 42)) + (0 if is_train else 17), + inference=False, + **kwargs, + ) + + +def _copy_vocab_assets(source_vocab_dir: Path, output_vocab_dir: Path, assay_vocab_to_copy: Path) -> None: + output_vocab_dir.mkdir(parents=True, exist_ok=True) + for file in source_vocab_dir.glob("*"): + if file.is_file() and file.name != "assay_vocab.json": + shutil.copy2(file, output_vocab_dir / file.name) + shutil.copy2(assay_vocab_to_copy, output_vocab_dir / "assay_vocab.json") + + +def _save_plain_weights(lightning_ckpt_path: Path, output_model_path: Path) -> None: + ckpt = torch.load(lightning_ckpt_path, map_location="cpu") + state_dict = ckpt["state_dict"] + cleaned = {k[len("model.") :]: v for k, v in state_dict.items() if k.startswith("model.")} + torch.save(cleaned, output_model_path) + + +def _save_dataset_filter_metadata(output_dir: Path, train_dataset, val_dataset=None) -> None: + summary = { + "train": getattr(train_dataset, "filter_metadata", None), + "val": getattr(val_dataset, "filter_metadata", None) if val_dataset is not None else None, + } + if summary["train"] is None and summary["val"] is None: + return + + with (output_dir / "data_filtering_summary.json").open("w") as f: + json.dump(summary, f, indent=2) + f.write("\n") + + +def _prepare_source_dirs(cfg: dict) -> tuple[Path, Path, Path]: + resume_dir = cfg.get("resume_artifact_dir") + if resume_dir: + source_dir = Path(resume_dir) + else: + source_dir = Path(cfg["checkpoint_dir"]) + + config_path = source_dir / "config.json" + weights_path = source_dir / "model_weights.pt" + vocab_dir = source_dir / "vocabs" + + if not config_path.exists() or not weights_path.exists() or not vocab_dir.exists(): + raise FileNotFoundError(f"Invalid artifact directory: {source_dir}") + + return config_path, weights_path, vocab_dir + + +def run_train_from_dict(cfg: dict) -> dict: + pl.seed_everything(int(cfg.get("seed", 42)), workers=True) + + config_path, weights_path, vocab_dir = _prepare_source_dirs(cfg) + + with config_path.open() as f: + config_json = json.load(f) + model_json = config_json["model"] + + cfg = dict(cfg) + cfg["source_vocab_dir"] = str(vocab_dir) + + data_config = _sanitize_data_config(model_json["data_config"], cfg) + model_config = _sanitize_model_config(model_json["model_config"]) + loss_config = _sanitize_loss_config(model_json["loss_config"], cfg) + + obs_assay_col = cfg.get("obs_assay_col", "assay") + aux_vocab = open_vocabs(str(vocab_dir), cols_to_load=None) + if obs_assay_col not in aux_vocab: + raise ValueError(f"Aux key '{obs_assay_col}' not found in {vocab_dir}") + + source_assay_vocab = aux_vocab[obs_assay_col] + expanded_path = cfg.get("expanded_assay_vocab") + if expanded_path: + with Path(expanded_path).open() as f: + target_assay_vocab = json.load(f) + assay_vocab_path_for_output = Path(expanded_path) + else: + target_assay_vocab = source_assay_vocab + assay_vocab_path_for_output = vocab_dir / "assay_vocab.json" + # Training aux tokenization updated with new assay vocab + aux_vocab[obs_assay_col] = target_assay_vocab + + emb_files = [str(vocab_dir / file_name) for file_name in data_config.esm2_mappings] + gene_vocab, emb_matrix = construct_gene_embeddings(emb_files, data_config.special_tokens) + emb_matrix = torch.tensor(emb_matrix) + + base_model = Transcriptformer( + data_config=data_config, + model_config=model_config, + loss_config=loss_config, + inference_config=None, + gene_vocab_dict=gene_vocab, + aux_vocab_dict=aux_vocab, + emb_matrix=emb_matrix, + ) + + state_dict = torch.load(weights_path, map_location="cpu", weights_only=True) + aux_key = f"aux_embeddings.{obs_assay_col}.weight" + if aux_key not in state_dict: + raise KeyError(f"Missing checkpoint key: {aux_key}") + + init_map = _parse_init_map(cfg.get("assay_init_map", [])) + state_dict[aux_key] = build_expanded_assay_embedding_weight( + old_weight=state_dict[aux_key], + old_vocab=source_assay_vocab, + new_vocab=target_assay_vocab, + init_map=init_map, + cfg=AssayInitConfig(default_source=cfg.get("init_default_source", "unknown"), mean_pool_fallback=True), + ) + + base_model.load_state_dict(state_dict, strict=False) + + apply_freeze_policy( + base_model, + freeze_transformer=bool(cfg.get("freeze_transformer", False)), + freeze_gene_embeddings=bool(cfg.get("freeze_gene_embeddings", False)), + freeze_count_head=bool(cfg.get("freeze_count_head", False)), + freeze_gene_head=bool(cfg.get("freeze_gene_head", False)), + train_aux_only=bool(cfg.get("train_aux_only", False)), + unfreeze_last_n_transformer_blocks=int(cfg.get("unfreeze_last_n_transformer_blocks", 0)), + ) + + trainable, total = count_trainable_parameters(base_model) + + train_dataset = _build_dataset( + files=list(cfg["train_files"]), + cfg=cfg, + data_config=data_config, + gene_vocab=gene_vocab, + aux_vocab=aux_vocab, + seq_len=model_config.seq_len, + is_train=True, + ) + val_files = list(cfg.get("val_files", [])) + val_dataset = ( + _build_dataset( + files=val_files, + cfg=cfg, + data_config=data_config, + gene_vocab=gene_vocab, + aux_vocab=aux_vocab, + seq_len=model_config.seq_len, + is_train=False, + ) + if val_files + else None + ) + + batch_size = int(cfg.get("batch_size", 8)) + num_workers = int(cfg.get("num_workers", 4)) + use_file_aware_batching = bool(cfg.get("enable_file_aware_batching", True)) + use_custom_batch_sampler = isinstance(train_dataset, AnnDatasetOOM) and use_file_aware_batching + if use_custom_batch_sampler: + train_loader = DataLoader( + train_dataset, + batch_sampler=FileAwareBatchSampler( + train_dataset, + batch_size=batch_size, + shuffle=True, + batches_per_file=int(cfg.get("oom_batches_per_file", 1)), + seed=int(cfg.get("seed", 42)), + ), + num_workers=num_workers, + pin_memory=True, + collate_fn=train_dataset.collate_fn, + ) + else: + train_loader = DataLoader( + train_dataset, + batch_size=batch_size, + # Keep shuffling enabled here; Lightning/DistributedSampler will own rank-aware shuffling in DDP. + shuffle=True, + num_workers=num_workers, + pin_memory=True, + collate_fn=train_dataset.collate_fn, + ) + val_loader = None + if val_dataset is not None: + if isinstance(val_dataset, AnnDatasetOOM) and use_file_aware_batching: + val_loader = DataLoader( + val_dataset, + batch_sampler=FileAwareBatchSampler( + val_dataset, + batch_size=batch_size, + shuffle=False, + batches_per_file=1, + seed=int(cfg.get("seed", 42)) + 17, + ), + num_workers=num_workers, + pin_memory=True, + collate_fn=val_dataset.collate_fn, + ) + else: + val_loader = DataLoader( + val_dataset, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=True, + collate_fn=val_dataset.collate_fn, + ) + + lit_model = TranscriptformerTrainModule( + model=base_model, + lr=float(cfg.get("lr", 5.5e-5)), + weight_decay=float(cfg.get("weight_decay", 0.05)), + beta1=float(cfg.get("adam_beta1", 0.9)), + beta2=float(cfg.get("adam_beta2", 0.95)), + eps=float(cfg.get("adam_eps", 1e-8)), + warmup_ratio=float(cfg.get("warmup_ratio", 0.1)), + min_lr_ratio=float(cfg.get("min_lr_ratio", 0.1)), + shuffle_expressed_each_batch=bool(cfg.get("shuffle_expressed_each_batch", False)), + ) + + + output_dir = Path(cfg["output_dir"]) + output_dir.mkdir(parents=True, exist_ok=True) + + _save_dataset_filter_metadata(output_dir, train_dataset, val_dataset) + + # Save config and train_args BEFORE training loop + out_config = dict(config_json) + out_config["model"] = dict(model_json) + out_config["model"]["data_config"] = dict(out_config["model"]["data_config"]) + out_config["model"]["data_config"]["aux_cols"] = obs_assay_col + out_config["train_runtime"] = cfg + with (output_dir / "config.json").open("w") as f: + json.dump(out_config, f, indent=2) + f.write("\n") + + with (output_dir / "train_args.json").open("w") as f: + json.dump(cfg, f, indent=2) + f.write("\n") + + + # Checkpoint config + ckpt_cfg = cfg.get("checkpoint", {}) + save_top_k = int(ckpt_cfg.get("save_top_k", 1)) + save_last = bool(ckpt_cfg.get("save_last", True)) + + monitor_metric = "val/total_loss" if val_loader is not None else "train/total_loss" + filename_template = "epoch-{epoch:03d}-step-{step:08d}-val-{val/total_loss:.4f}" + if val_loader is None: + filename_template = "epoch-{epoch:03d}-step-{step:08d}-train-{train/total_loss:.4f}" + + callbacks: list[pl.Callback] = [] + ckpt_topk = None + ckpt_last = None + + # Top-K metric-based checkpoints + if save_top_k > 0: + ckpt_topk = ModelCheckpoint( + dirpath=str(output_dir / "lightning_ckpts"), + filename=filename_template, + monitor=monitor_metric, + mode="min", + save_top_k=save_top_k, + save_last=False, + auto_insert_metric_name=False, # avoid duplicate metric name and "=" sign + ) + callbacks.append(ckpt_topk) + + # Last checkpoint at each epoch end + if save_last: + ckpt_last = ModelCheckpoint( + dirpath=str(output_dir / "lightning_ckpts"), + filename="last", + save_top_k=0, + save_last=True, + every_n_epochs=1, + ) + callbacks.append(ckpt_last) + + device_preference = cfg.get("device", cfg.get("accelerator", "auto")) + num_gpus = cfg.get("num_gpus", cfg.get("devices", 1)) + accelerator, devices = resolve_accelerator_and_devices(device_preference, num_gpus) + num_nodes = int(cfg.get("num_nodes", 1)) + strategy = "ddp" if (int(devices) > 1 or num_nodes > 1) else "auto" + + # Gradient accumulation and clipping + grad_clip = float(cfg.get("gradient_clip_val", 1.0)) + grad_accum = int(cfg.get("accumulate_grad_batches", 1)) + + # CSV Logger for live metrics + from pytorch_lightning.loggers import CSVLogger + csv_logger = CSVLogger(save_dir=str(output_dir), name="csv_logs") + + live_cfg = cfg.get("live_metrics", {}) + live_metrics_cb = LiveMetricsSummaryCallback( + output_dir=output_dir, + enabled=bool(live_cfg.get("enabled", True)), + every_n_epochs=int(live_cfg.get("every_n_epochs", 1)), + ) + callbacks.append(live_metrics_cb) + + trainer = pl.Trainer( + accelerator=accelerator, + devices=devices, + num_nodes=num_nodes, + precision=cfg.get("precision", "16-mixed"), + max_epochs=int(cfg.get("max_epochs", 5)), + log_every_n_steps=int(cfg.get("log_every_n_steps", 10)), + callbacks=callbacks, + strategy=strategy, + gradient_clip_val=grad_clip, + accumulate_grad_batches=grad_accum, + logger=csv_logger, + use_distributed_sampler=not use_custom_batch_sampler, + ) + + resume_mode = cfg.get("resume_mode", "weights") + resume_ckpt = None + if cfg.get("resume_artifact_dir") and resume_mode == "lightning": + resume_ckpt = Path(cfg["resume_artifact_dir"]) / "lightning_ckpts" / "last.ckpt" + if not resume_ckpt.exists(): + raise FileNotFoundError(f"Missing resume checkpoint: {resume_ckpt}") + + trainer.fit( + lit_model, + train_dataloaders=train_loader, + val_dataloaders=val_loader, + ckpt_path=str(resume_ckpt) if resume_ckpt else None, + ) + + best_ckpt_path = "" + if ckpt_topk is not None and ckpt_topk.best_model_path: + best_ckpt_path = ckpt_topk.best_model_path + elif ckpt_last is not None: + best_ckpt_path = str(Path(ckpt_last.dirpath) / "last.ckpt") + + best_ckpt = Path(best_ckpt_path) + if not best_ckpt.exists(): + raise RuntimeError("No checkpoint generated") + + output_weights = output_dir / "model_weights.pt" + _save_plain_weights(best_ckpt, output_weights) + + output_vocab_dir = output_dir / "vocabs" + _copy_vocab_assets(vocab_dir, output_vocab_dir, assay_vocab_path_for_output) + + return { + "output_dir": str(output_dir), + "model_weights": str(output_weights), + "trainable_params": trainable, + "total_params": total, + } + + +def setup_runtime_for_training() -> None: + torch._dynamo.config.optimize_ddp = False + os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") diff --git a/src/transcriptformer/utils/device.py b/src/transcriptformer/utils/device.py new file mode 100644 index 0000000..f3a10da --- /dev/null +++ b/src/transcriptformer/utils/device.py @@ -0,0 +1,107 @@ +import logging + +import torch + + +def _normalize_device_preference(device_preference: str | None) -> str: + if device_preference is None: + return "auto" + value = str(device_preference).strip().lower() + if value == "gpu": + return "cuda" + if value in {"auto", "cpu", "cuda", "mps"}: + return value + raise ValueError(f"Unsupported device preference '{device_preference}'. Use one of: auto, cpu, cuda, mps") + + +def _parse_num_gpus(num_gpus: int | str | None) -> int: + if num_gpus is None: + return 1 + if isinstance(num_gpus, int): + return num_gpus + + value = str(num_gpus).strip().lower() + if value in {"-1", "all", "auto"}: + return -1 + return int(value) + + +def resolve_checkpoint_map_location(device_preference: str | None = "auto") -> str: + """Resolve checkpoint map_location from a device preference.""" + device = _normalize_device_preference(device_preference) + if device == "cpu": + return "cpu" + if device == "cuda": + return "cuda" if torch.cuda.is_available() else "cpu" + if device == "mps": + return "mps" if torch.backends.mps.is_available() else "cpu" + + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return "cpu" + + +def resolve_accelerator_and_devices( + device_preference: str | None = "auto", + num_gpus: int | str | None = 1, +) -> tuple[str, int]: + """Resolve Lightning accelerator/devices from the shared device policy.""" + device = _normalize_device_preference(device_preference) + requested = _parse_num_gpus(num_gpus) + + if device == "cpu": + return "cpu", 1 + + if device == "mps": + if not torch.backends.mps.is_available(): + raise RuntimeError("MPS requested but not available") + return "mps", 1 + + if device == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA requested but not available") + available = torch.cuda.device_count() + if requested == -1: + return "gpu", available + if requested > 1: + if available < requested: + logging.warning( + "Requested %s CUDA devices but only %s available. Using %s.", + requested, + available, + available, + ) + return "gpu", min(requested, available) + return "gpu", 1 + + # auto + if requested == -1: + if torch.cuda.is_available(): + return "gpu", torch.cuda.device_count() + if torch.backends.mps.is_available(): + return "mps", 1 + return "cpu", 1 + + if requested > 1: + available = torch.cuda.device_count() if torch.cuda.is_available() else 0 + if available < requested: + logging.warning( + "Requested %s GPUs but only %s available. Using %s GPU(s).", + requested, + available, + available if available > 0 else 1, + ) + if available > 0: + return "gpu", available + if torch.backends.mps.is_available(): + return "mps", 1 + return "cpu", 1 + return "gpu", requested + + if torch.cuda.is_available(): + return "gpu", 1 + if torch.backends.mps.is_available(): + return "mps", 1 + return "cpu", 1 \ No newline at end of file diff --git a/test/test_cli.py b/test/test_cli.py index 428f73b..2248dea 100644 --- a/test/test_cli.py +++ b/test/test_cli.py @@ -1,5 +1,6 @@ """Tests for the TranscriptFormer CLI module.""" +import argparse import sys from unittest import mock @@ -7,9 +8,9 @@ from transcriptformer.cli import ( main, - run_download_cli, - setup_download_parser, + run_train_cli, setup_inference_parser, + setup_train_parser, ) @@ -18,13 +19,11 @@ class TestCLIMain: def test_main_no_args(self, monkeypatch, capsys): """Test CLI with no arguments prints help and exits.""" - # Mock sys.argv and sys.exit monkeypatch.setattr(sys, "argv", ["transcriptformer"]) with mock.patch("sys.exit") as mock_exit: main() mock_exit.assert_called_once_with(1) - # Check help was printed captured = capsys.readouterr() assert "usage: " in captured.out assert "TranscriptFormer command-line interface" in captured.out @@ -63,42 +62,106 @@ def test_inference_command(self, mock_run_inference, monkeypatch): mock_run_inference.assert_called_once() -class TestDownloadCommand: - """Tests for the download command.""" +class TestTrainCommand: + """Tests for the train command.""" - @mock.patch("transcriptformer.cli.run_download_cli") - def test_download_command(self, mock_run_download, monkeypatch): - """Test that download command runs with required arguments.""" - monkeypatch.setattr(sys, "argv", ["transcriptformer", "download", "tf-sapiens"]) + @mock.patch("transcriptformer.cli.run_train_cli") + def test_train_command(self, mock_run_train, monkeypatch): + """Test that train command runs with required arguments.""" + monkeypatch.setattr( + sys, + "argv", + [ + "transcriptformer", + "train", + "--checkpoint-dir", + "/path/to/checkpoint", + "--output-dir", + "/path/to/output", + "--train-file", + "/path/to/train.h5ad", + ], + ) main() - mock_run_download.assert_called_once() - - @mock.patch("transcriptformer.cli.download_artifacts.download_and_extract") - def test_run_download_cli_single_model(self, mock_download, monkeypatch): - """Test run_download_cli with a single model.""" - args = mock.MagicMock() - args.model = "tf-sapiens" - args.checkpoint_dir = "./checkpoints" - - run_download_cli(args) - mock_download.assert_called_once_with("tf_sapiens", "./checkpoints") + mock_run_train.assert_called_once() - @mock.patch("transcriptformer.cli.download_artifacts.download_and_extract") - def test_run_download_cli_all_models(self, mock_download, monkeypatch): - """Test run_download_cli with 'all' option.""" + @mock.patch("transcriptformer.cli.run_train_from_dict") + @mock.patch("transcriptformer.cli.setup_runtime_for_training") + def test_run_train_cli(self, mock_setup_runtime, mock_run_train_from_dict): + """Test run_train_cli parameter mapping.""" args = mock.MagicMock() - args.model = "all" - args.checkpoint_dir = "./checkpoints" - - run_download_cli(args) - assert mock_download.call_count == 4 - - # Check all models were downloaded - mock_download.assert_any_call("tf_sapiens", "./checkpoints") - mock_download.assert_any_call("tf_exemplar", "./checkpoints") - mock_download.assert_any_call("tf_metazoa", "./checkpoints") - mock_download.assert_any_call("all_embeddings", "./checkpoints") + args.checkpoint_dir = "/path/to/checkpoint" + args.resume_artifact_dir = None + args.resume_mode = "weights" + args.output_dir = "/path/to/output" + args.train_file = ["/path/to/train.h5ad"] + args.val_file = [] + args.expanded_assay_vocab = None + args.obs_assay_col = "assay" + args.gene_col_name = "ensembl_id" + args.filter_to_vocabs = True + args.filter_outliers = 0.0 + args.sort_genes = False + args.randomize_genes = False + args.min_expressed_genes = 0 + args.n_data_workers = 4 + args.batch_size = 2 + args.num_workers = 0 + args.max_epochs = 1 + args.precision = "32" + args.device = "cpu" + args.num_gpus = 1 + args.num_nodes = 1 + args.lr = 1e-4 + args.weight_decay = 0.0 + args.adam_beta1 = 0.9 + args.adam_beta2 = 0.95 + args.adam_eps = 1e-8 + args.warmup_ratio = 0.1 + args.min_lr_ratio = 0.1 + args.gene_id_loss_weight = 1.0 + args.softplus_approx = True + args.init_default_source = "unknown" + args.assay_init_map = [] + args.freeze_transformer = False + args.unfreeze_last_n_transformer_blocks = 0 + args.freeze_gene_embeddings = False + args.freeze_count_head = False + args.freeze_gene_head = False + args.train_aux_only = False + args.shuffle_expressed_each_batch = False + args.clip_counts = 30.0 + args.normalize_to_scale = 0.0 + args.use_raw = False + args.remove_duplicate_genes = False + args.use_oom_dataloader = False + args.file_aware_batching = True + args.oom_batches_per_file = 1 + args.seed = 42 + + # Mock return value + mock_run_train_from_dict.return_value = {"output_dir": "/path/to/output"} + + # Call the function + run_train_cli(args) + + # Verify setup was called + mock_setup_runtime.assert_called_once() + + # Verify run_train_from_dict was called with correct config + mock_run_train_from_dict.assert_called_once() + call_args = mock_run_train_from_dict.call_args[0][0] + assert call_args["checkpoint_dir"] == "/path/to/checkpoint" + assert call_args["output_dir"] == "/path/to/output" + assert call_args["train_files"] == ["/path/to/train.h5ad"] + assert call_args["data_config"]["gene_col_name"] == "ensembl_id" + assert call_args["device"] == "cpu" + assert call_args["num_gpus"] == 1 + assert call_args["unfreeze_last_n_transformer_blocks"] == 0 + assert call_args["enable_file_aware_batching"] is True + assert call_args["oom_batches_per_file"] == 1 + assert call_args["loss_config"]["gene_id_loss_weight"] == 1.0 class TestCLIParsers: @@ -112,14 +175,12 @@ def test_inference_parser_setup(self): setup_inference_parser(subparsers) - # Check required arguments are added subparsers.add_parser.assert_called_once_with( "inference", help="Run inference with a TranscriptFormer model", description="Run inference with a TranscriptFormer model on scRNA-seq data.", ) - # Check that required arguments are added with required=True parser.add_argument.assert_any_call( "--checkpoint-path", required=True, @@ -131,7 +192,6 @@ def test_inference_parser_setup(self): help="Path to input AnnData file to run inference on", ) - # Check that emb-type argument is added with correct choices parser.add_argument.assert_any_call( "--emb-type", default="cell", @@ -139,30 +199,39 @@ def test_inference_parser_setup(self): help="Type of embeddings to extract: 'cell' for mean-pooled cell embeddings or 'cge' for contextual gene embeddings (default: cell)", ) - def test_download_parser_setup(self): - """Test that download parser is set up correctly.""" + def test_train_parser_setup(self): + """Test that train parser is set up correctly.""" parser = mock.MagicMock() subparsers = mock.MagicMock() subparsers.add_parser.return_value = parser - setup_download_parser(subparsers) + setup_train_parser(subparsers) - # Check parser is created subparsers.add_parser.assert_called_once_with( - "download", - help="Download and extract TranscriptFormer model artifacts", - description="Download and extract TranscriptFormer model artifacts from a public S3 bucket.", + "train", + help="Train or continue training with expanded assay vocab", + description="Fine-tune TranscriptFormer with expanded assay tokens and optional freezing.", ) - # Check that model argument is added with choices parser.add_argument.assert_any_call( - "model", - choices=[ - "tf-sapiens", - "tf-exemplar", - "tf-metazoa", - "all", - "all-embeddings", - ], - help="Model to download ('all' for all models and embeddings, 'all-embeddings' for just embeddings)", + "--checkpoint-dir", + required=True, + help="Base artifact directory with config.json/model_weights.pt", + ) + parser.add_argument.assert_any_call( + "--output-dir", + required=True, + help="Output artifact directory", + ) + parser.add_argument.assert_any_call( + "--train-file", + action="append", + required=True, + help="Training .h5ad file (repeatable)", + ) + parser.add_argument.assert_any_call( + "--file-aware-batching", + action=argparse.BooleanOptionalAction, + default=True, + help="Keep OOM batches mostly file-local; disable to fall back to Lightning DistributedSampler batching", ) diff --git a/test/test_train_smoke.py b/test/test_train_smoke.py new file mode 100644 index 0000000..dbec5b0 --- /dev/null +++ b/test/test_train_smoke.py @@ -0,0 +1,109 @@ +"""Smoke tests for training module.""" + +import pytest +import torch + +from transcriptformer.data.dataclasses import BatchData, DataConfig, LossConfig, ModelConfig +from transcriptformer.model.model import Transcriptformer +from transcriptformer.train.engine import TranscriptformerTrainModule + + +def _tiny_model() -> Transcriptformer: + gene_vocab = { + "unknown": 0, + "[PAD]": 1, + "[START]": 2, + "[END]": 3, + "[RD]": 4, + "[CELL]": 5, + "[MASK]": 6, + "g1": 7, + "g2": 8, + "g3": 9, + } + aux_vocab = {"assay": {"unknown": 0, "new_assay": 1}} + + data_config = DataConfig( + aux_vocab_path=".", + pin_memory=False, + aux_cols=["assay"], + gene_col_name="ensembl_id", + clip_counts=30, + filter_to_vocabs=True, + filter_outliers=0.0, + pad_zeros=True, + normalize_to_scale=0, + n_data_workers=0, + sort_genes=False, + randomize_genes=False, + min_expressed_genes=0, + gene_pad_token="[PAD]", + aux_pad_token="unknown", + ) + + model_config = ModelConfig( + log_counts_eps=1e-6, + num_heads=2, + num_layers=1, + model_dim=16, + embed_dim=8, + dropout=0.1, + activation="gelu", + attn_bias=False, + fw_bias=False, + mu_link_fn="softplus", + softcap=0, + seq_len=3, + aux_len=1, + block_len=2, + compile_block_mask=False, + ) + + loss_config = LossConfig(gene_id_loss_weight=1.0, softplus_approx=True) + + emb_matrix = torch.randn(len(gene_vocab), model_config.embed_dim) + + return Transcriptformer( + data_config=data_config, + model_config=model_config, + loss_config=loss_config, + gene_vocab_dict=gene_vocab, + aux_vocab_dict=aux_vocab, + emb_matrix=emb_matrix, + ) + + +def test_train_smoke_forward_backward(): + """Smoke test for training forward/backward (skipped on CPU due to FlexAttention limitation).""" + # FlexAttention does not support backward on CPU, so skip this test if no GPU available + if not torch.cuda.is_available(): + pytest.skip("FlexAttention backward not supported on CPU; GPU required for this test") + + device = torch.device("cuda") + model = _tiny_model() + model = model.to(device) + + module = TranscriptformerTrainModule( + model=model, + lr=1e-4, + weight_decay=0.0, + beta1=0.9, + beta2=0.95, + eps=1e-8, + warmup_ratio=0.1, + min_lr_ratio=0.1, + shuffle_expressed_each_batch=False, + ) + + batch = BatchData( + gene_counts=torch.tensor([[5.0, 3.0, 2.0], [2.0, 1.0, 4.0]], dtype=torch.float32, device=device), + gene_token_indices=torch.tensor([[7, 8, 9], [8, 7, 9]], dtype=torch.int64, device=device), + aux_token_indices=torch.tensor([[0], [1]], dtype=torch.int64, device=device), + ) + + loss = module.training_step(batch, 0) + assert torch.isfinite(loss).item() + loss.backward() + + grads = [p.grad for p in model.parameters() if p.requires_grad] + assert any(g is not None for g in grads)