scdynomics.Masked_Pretrainer
- class scdynomics.Masked_Pretrainer(*args: Any, **kwargs: Any)
Bases:
LightningModulePretrainer module with input masking and reconstruction for self-supervised learning.
This core PyTorch Lightning module implements input masking and relies on the model to reconstruct either full floating-point values or binned values, driving the foundation model to learn robust gene representations and multiomics dynamics.
- __init__(attention: str = 'linformer', train_mode: str = 'float', species: str = 'mm10', multispecies: bool = False, token_dict: dict = None, input_drop_prob: float = 0.0, input_mod: int = 2, value_bins: int = 15, binning_factor: int = 1, embed_dim: int = 1, dim_key: int = None, dim_value: int = None, attn_dropout: float = 0.0, causal: bool = False, projection_len: int = None, num_heads: int = 1, dim_feedforward: int = 512, dropout_prob: float = 0.0, bias: bool = True, glu: bool = False, layer_norm_eps: float = 1e-05, use_rotary_embeddings: bool = False, num_layers: int = 1, pool_method: str = None, pool_out_size: int = 1, loss_reduction: str = 'mean', optimizer: str = 'adamw', weight_decay: float = 0.01, betas=(0.9, 0.999), momentum: float = 0.9, learning_rate: float = 0.0001, sch_warmpup_epochs: int = None, sch_warmpup_factor: float = 1, sch_T_0: int = 20, sch_T_mult: int = 2, min_lr: float = 2e-06, total_steps: int = 0, steps_per_epoch: int = 0, record_grad_norm: bool = False, **kwargs)
- Parameters:
- attention:
str(default:"linformer") Type of attention mechanism to use (“linformer”, “tf”, “scaled_dot_product”).
- train_mode:
str(default:"float") Training reconstruction mode (“float” for continuous, “bin” for cross-entropy).
- species:
str(default:'mm10') Species token dictionary to train on (e.g., ‘mm10’, ‘hg38’).
- multispecies:
bool(default:False) Whether to train on multi-species tokens.
- token_dict:
dict(default:None) Dictionary containing token mappings.
- input_drop_prob:
float(default:0.0) Probability of dropping input elements for masking.
- input_mod:
int(default:2) Number of input modalities.
- value_bins:
int(default:15) Number of value bins for the embedding.
- binning_factor:
int(default:1) Factor for value binning.
- embed_dim:
int(default:1) Dimension of the embedding.
- dim_key:
int(default:None) Dimension size of the key projections in attention.
- dim_value:
int(default:None) Dimension size of the value projections in attention.
- attn_dropout:
float(default:0.0) Dropout probability applied specifically to attention weights.
- causal:
bool(default:False) If True, applies a causal mask to the attention layer.
- projection_len:
int(default:None) Projection length used in Linformer attention.
- num_heads:
int(default:1) Number of attention heads.
- dim_feedforward:
int(default:512) Hidden dimension size of the feedforward network.
- dropout_prob:
float(default:0.0) Dropout probability for the feedforward network.
- bias:
bool(default:True) If True, enables bias in linear layers.
- glu:
bool(default:False) If True, utilizes Gated Linear Units in the feedforward network.
- layer_norm_eps:
float(default:1e-5) Epsilon value for Layer Normalization.
- use_rotary_embeddings:
bool(default:False) If True, utilizes Rotary Position Embeddings (RoPE).
- num_layers:
int(default:1) Number of consecutive encoder blocks.
- pool_method:
str(default:None) Method for pooling sequence embeddings (“avg”, “max”, “dense”).
- pool_out_size:
int(default:1) Output size for the pooling layer.
- loss_reduction:
str(default:"mean") Specifies the reduction to apply to the output loss.
- optimizer:
str(default:"adamw") Optimizer to use (e.g., “adamw”).
- weight_decay:
float(default:0.01) Weight decay (L2 penalty) for the optimizer.
- betas:
tuple(default:(0.9, 0.999)) Beta coefficients for AdamW.
- momentum:
float(default:0.9) Momentum factor for applicable optimizers.
- learning_rate:
float(default:1e-4) Base learning rate for training.
- sch_warmpup_epochs:
int(default:None) Number of epochs to linearly warm up the learning rate.
- sch_warmpup_factor:
float(default:1.0) Multiplier for the learning rate warmup.
- sch_T_0:
int(default:20) Number of iterations for the first restart in CosineAnnealingWarmRestarts.
- sch_T_mult:
int(default:2) A factor increases Ti after a restart in CosineAnnealingWarmRestarts.
- min_lr:
float(default:2e-6) Minimum learning rate boundary for the scheduler.
- total_steps:
int(default:0) Total number of training steps for OneCycleLR.
- steps_per_epoch:
int(default:0) Steps per epoch for scheduling calculations.
- record_grad_norm:
bool(default:False) If True, calculates and logs the gradient norm during training.
- kwargs:
dict Additional keyword arguments.
- attention:
- cal_grad_norm() float
Calculate the L2 gradient norm across all model parameters.
- Returns:
float The scalar value of the calculated gradient norm.
- Returns:
- configure_optimizers() dict
Create the optimizer and the training learning rate schedule.
- Returns:
dict A dictionary containing the instantiated optimizer and learning rate scheduler configurations compatible with PyTorch Lightning.
- Returns:
- forward(x, masking: bool = False, **kwargs) torch.Tensor
- partition(embed_mod: list = None, query_genes: pandas.Index = None)
Partition the pretrainer model based on the specified modalities and genes.
- Parameters:
- embed_mod:
list(default: None) The embedding modalities to be kept. Example of only using the second modality:
[1]. IfNone, no partitioning will be done.- query_genes:
pd.Index(default: None) The genes invovled in the fine-tuning task. If
None, no partitioning will be done.
- embed_mod:
- Returns:
pd.Index The list of features (genes) to be kept.
- test_step(batch, batch_idx: int = None) None
Defines the test loop logic
- to_dev(device) Masked_Pretrainer
Move explicit model components directly to the specified computing device.
- Parameters:
- device:
torch.deviceorstr The target device representation (e.g., ‘cuda:0’, ‘cpu’).
- device:
- Returns:
Masked_Pretrainer The instance of the model with parameters migrated to the device.
- property token_tensor: torch.Tensor
Gene token tensor with int type
- training_step(batch, batch_idx: int = None) torch.Tensor
Defines the training loop logic
- validation_step(batch, batch_idx: int = None) None
Defines the validation loop logic