scdynomics.FT_MLP_Classifier
- class scdynomics.FT_MLP_Classifier(*args: Any, **kwargs: Any)
Bases:
LightningModuleSupervised Classification Tuner Using MLP Modules over a foundation model.
Connects an underlying FT_Encoder to a flexible Multi-Layer Perceptron (MLP) decision head. Integrates capabilities for robust downstream cellular classification tasks and includes methods for interpretability analysis.
- __init__(base_model=None, unfreeze_embedder: bool = True, unfreeze_embedder_target: str = 'sgt', lora_type: str = 'adapter', shortcut_type: str = None, pre_block_num: int = None, lora_latent_dim: int = 512, lora_dropout: float = 0.0, lora_bias: bool = False, lora_activation: str = 'relu', loss_reduction: str = 'mean', pool_method: str = None, pool_attn_hidden: int = 32, pool_out_size: int = 1, n_decision_layers: int = 3, decision_res_connect: bool = False, decision_dropout: float = 0.1, decision_hidden_dim: int = 512, n_classes: int = 2, optimizer: str = 'adamw', weight_decay: float = 0.001, betas=(0.9, 0.98), momentum: float = 0.9, learning_rate: float = 0.0001, scheduler: str = 'cosine', 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 = 1, steps_per_epoch: int = 1, record_grad_norm: bool = False, **kwargs)
- Parameters:
- base_model:
Masked_Pretrainer(default:None) The pre-trained foundation model containing tokens and embedding logic.
- unfreeze_embedder:
bool(default:True) Whether to unfreeze the initial embedding layers.
- unfreeze_embedder_target:
str(default:"sgt") Which sub-embedding to unfreeze:
"sgt"(gene-token / position embedding),"val"(value-bin embedding), or"all".- lora_type:
str(default:"adapter") The PEFT architecture to apply within the FT_Encoder (“adapter”, “lora”).
- shortcut_type:
str(default:None) Type of shortcut connection.
- pre_block_num:
int(default:None) Number of initial blocks configured.
- lora_latent_dim:
int(default:512) The rank dimension for the PEFT matrices.
- lora_dropout:
float(default:0.0) Dropout rate for LoRA.
- lora_bias:
bool(default:False) Whether to enable bias training in LoRA.
- lora_activation:
str(default:'relu') Activation function for the LoRA adaptation.
- loss_reduction:
str(default:"mean") Defines the reduction method for CrossEntropyLoss.
- pool_method:
str(default:None) Pooling method for sequence representations.
- pool_attn_hidden:
int(default:32) Internal MLP dimension for attention pooling.
- pool_out_size:
int(default:1) Size of the pooled output.
- n_decision_layers:
int(default:3) Total number of sequential MLP blocks in the classification head.
- decision_res_connect:
bool(default:False) If True, adds residual connections between internal MLP layers.
- decision_dropout:
float(default:0.1) Dropout rate applied within the decision layers.
- decision_hidden_dim:
int(default:512) Feature dimension size of the internal decision MLP layers.
- n_classes:
int(default:2) The number of output classes for the prediction task.
- optimizer:
str(default:"adamw") Optimizer to use for training.
- weight_decay:
float(default:0.001) Weight decay for the optimizer.
- betas:
tuple(default:(0.9, 0.98)) Beta parameters for AdamW.
- momentum:
float(default:0.9) Momentum for optimization.
- learning_rate:
float(default:1e-4) Learning rate for the optimizer.
- scheduler:
str(default:"cosine") Type of learning rate scheduler to use.
- sch_warmpup_epochs:
int(default:None) Epochs to warm up the learning rate.
- sch_warmpup_factor:
float(default:1.0) Factor for warming up the learning rate.
- sch_T_0:
int(default:20) Number of iterations for the first restart in scheduler.
- sch_T_mult:
int(default:2) Multiplier for restarts in scheduler.
- min_lr:
float(default:2e-6) Minimum learning rate.
- total_steps:
int(default:1) Total steps for scheduling.
- steps_per_epoch:
int(default:1) Steps per epoch for scheduling.
- record_grad_norm:
bool(default:False) If True, logs the gradient norm during training.
- kwargs:
dict Additional optimization and scheduling arguments.
- base_model:
- cal_grad_norm() float
Calculate the gradient norm across all model parameters.
- Returns:
float The scalar value of the calculated gradient norm.
- Returns:
- configure_optimizers()
Create the optimizer and the training schedule.
- Returns:
dict A dictionary containing the optimizer and scheduler configuration.
- Returns:
- explain(dataset=None, exp_batch_size: int = 5, sample_limit: int = None, sample_method: str = 'top', n_dataloader_workers: int = 10, device: str = 'cuda', precision: str = 'medium', explain_method: str = 'gausslegendre', explain_step: int = 50, verbose: bool = False)
Explain a binary classifier model.
- Parameters:
- dataset:
Dataset(default:None) The evaluation dataset used for explanation.
- exp_batch_size:
int(default:5) Batch size to use during explanation inference.
- sample_limit:
int(default:None) Maximum number of samples to evaluate per class.
- sample_method:
str(default:'top') Method used to subsample the dataset for explanation.
- n_dataloader_workers:
int(default:10) Number of workers used in the explanation DataLoader.
- device:
str(default:"cuda") The computing device to utilize.
- precision:
str(default:"medium") Tensor precision setting.
- explain_method:
str(default:'gausslegendre') The approximation method for Integrated Gradients.
- explain_step:
int(default:50) The number of steps used in the IG approximation.
- verbose:
bool(default:False) If True, prints progress and duration logs.
- dataset:
- Returns:
pd.DataFrame A Pandas DataFrame containing the calculated feature attribution scores and standard deviations per class, indexed by features.
- explaining(mode: bool = True)
Set model to explaining mode.
- Parameters:
- mode:
bool(default:True) If True, activates explaining mode and unfreezes the base model.
- mode:
- forward(x, **kwargs)
- get_confmat(dataloader)
Make a confusion matrix based on the given dataloader.
- Parameters:
- dataloader:
DataLoader The PyTorch DataLoader providing the evaluation samples.
- dataloader:
- Returns:
ConfusionMatrix A torchmetrics ConfusionMatrix object populated with the predictions against true labels.
- get_init_embedding(x)
Retrieve the initial token embeddings prior to processing through the encoder.
- Parameters:
- x:
torch.Tensor The input tensor to embed.
- x:
- Returns:
torch.Tensor The detached embedded tensor.
- predict(x, y=None, return_np: bool = True)
Predict the class probabilities of the input data.
- Parameters:
- x:
torch.Tensor Input data tensor.
- y:
torch.Tensor(default:None) True labels (unused in prediction, kept for API compatibility).
- return_np:
bool(default:True) If True, returns a numpy array instead of a PyTorch tensor.
- x:
- Returns:
np.ndarrayortorch.Tensor The softmax normalized prediction probabilities.
- test_step(batch, batch_idx: int = None)
Testing loop logic. Metrics: loss, accuracy, f1, auroc
- Parameters:
- batch:
tuple A tuple containing (x, y).
- batch_idx:
int(default:None) The index of the current batch.
- batch:
- training_step(batch, batch_idx: int = None)
Defines the training loop logic.
- Parameters:
- batch:
tuple A tuple containing (x, y) where x is the input data and y is the target labels.
- batch_idx:
int(default:None) The index of the current batch.
- batch:
- Returns:
torch.Tensor The computed CrossEntropy loss for the batch.
- validation_step(batch, batch_idx: int = None)
Validation loop logic. Metrics: loss, accuracy, f1, auroc
- Parameters:
- batch:
tuple A tuple containing (x, y).
- batch_idx:
int(default:None) The index of the current batch.
- batch: