torchcell.trainers.int_transformer_cell.RegressionTask

class RegressionTask(model: Module, cell_graph: Tensor, optimizer_config: dict[str, Any], lr_scheduler_config: dict[str, Any], batch_size: int | None = None, clip_grad_norm: bool = False, clip_grad_norm_max_norm: float = 0.1, plot_sample_ceiling: int = 1000, plot_every_n_epochs: int = 10, plot_transformer_diagnostics_every_n_epochs: int = 10, plot_edge_recovery_every_n_epochs: int = 10, loss_func: Module | None = None, grad_accumulation_schedule: dict[int, int] | None = None, device: str = 'cuda', inverse_transform: Module | None = None, execution_mode: str = 'training')

Bases: LightningModule

Lightning module training the transformer cell model on gene interactions.

reset_edge_recovery_accumulators() → None

Reset accumulators for edge recovery metrics.

reset_attention_diagnostics() → None

Reset attention diagnostic accumulators.

forward(batch: HeteroData, return_attention: bool = False) → Any

Run a forward pass, optionally returning attention diagnostics.

training_step(batch: HeteroData, batch_idx: int) → Tensor

Run a manual-optimization training step with gradient accumulation.

validation_step(batch: HeteroData, batch_idx: int) → Tensor

Run the validation shared step and periodically free CUDA cache.

test_step(batch: HeteroData, batch_idx: int) → Tensor

Run the shared step for the test stage and return the loss.

on_train_epoch_end() → None

Log metrics, plot samples, step the scheduler, and clear CUDA cache.

on_train_epoch_start() → None

Update gradient accumulation steps and clear training sample buffers.

on_validation_epoch_start() → None

Free CUDA memory, reset diagnostics, and clear validation sample buffers.

on_test_epoch_start() → None

Clear test sample buffers at the start of the test epoch.

on_validation_epoch_end() → None

Log and reset validation metrics and plot validation samples periodically.

Under manual optimization Lightning never steps a scheduler, so a ReduceLROnPlateau scheduler is stepped here, once per (non-sanity) validation epoch, on the value of its monitor val/gene_interaction/MSE computed just above. Validation runs inside the training epoch, before on_train_epoch_end, so this is the point where the monitor is fresh. Without a validation loop the plateau scheduler is never stepped and the learning rate stays fixed.

on_test_epoch_end() → None

Log and reset test metrics and plot test samples.

configure_optimizers() → Any

Build the optimizer and learning-rate scheduler from hparams config.