torchcell.trainers.int_hetero_cell.RegressionTask

class RegressionTask(model: Module, cell_graph: Tensor, optimizer_config: dict[str, Any], lr_scheduler_config: dict[str, Any] | None, 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, loss_func: Module | None = None, grad_accumulation_schedule: dict[Any, int] | None = None, device: str = 'cuda', inverse_transform: Module | None = None, execution_mode: str = 'training')

Bases: LightningModule

Lightning task training a model to predict gene interaction scores.

property hp: _HParams

Return self.hparams as a precisely typed view (runtime no-op cast).

forward(batch: HeteroData) → Any

Run the model on a batch and return predictions and any auxiliary outputs.

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

Run a manual-optimization training step and log loss and metrics.

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

Run a validation step, accumulating predictions and logging metrics.

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

Run a test step, accumulating predictions and logging metrics.

on_train_epoch_end() → None

Log epoch-level training metrics, plot, and step an epoch scheduler.

Under manual optimization Lightning steps no scheduler, so this hook steps a warmup or cosine scheduler once per epoch. A ReduceLROnPlateau scheduler is not stepped here: it needs its metric and is stepped in on_validation_epoch_end.

on_train_epoch_start() → None

Update gradient accumulation and reset training sample accumulators.

on_validation_epoch_start() → None

Reset validation sample accumulators at the start of the epoch.

on_test_epoch_start() → None

Reset test sample accumulators at the start of the epoch.

on_validation_epoch_end() → None

Log validation metrics, step a plateau scheduler, plot, and reset.

A ReduceLROnPlateau scheduler is stepped here, once per validation epoch outside the sanity check, on PLATEAU_MONITOR (val/gene_interaction/MSE) as just computed: validation runs inside the training epoch, so this is where the monitor is fresh, and an epoch without validation does not step it. A validation epoch whose collection received no rows is refused rather than stepped on NaN.

on_test_epoch_end() → None

Log test metrics, render plots, and reset accumulators.

configure_optimizers() → Optimizer | Sequence[Optimizer] | tuple[Sequence[Optimizer], Sequence[LRScheduler | ReduceLROnPlateau | LRSchedulerConfig]] | OptimizerConfig | OptimizerLRSchedulerConfig | Sequence[OptimizerConfig] | Sequence[OptimizerLRSchedulerConfig] | None

Build the optimizer and learning-rate scheduler from config.