torchcell.trainers.int_dango.RegressionTask

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

Bases: LightningModule

Lightning module training DANGO to regress genetic interaction scores.

forward(batch: Any) → tuple[Tensor, dict[str, Any]]

Forward pass through the model

Parameters:

batch – HeteroData batch containing perturbation information

Returns:

the model’s own outputs dict, unchanged (for Dango: integrated_embeddings, reconstructions, network_embeddings, initial_embeddings, interaction_scores), so one call feeds the loss, the metrics and every log of a step.

Return type:

Tuple of (predictions, outputs)

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

Run a training step and return the loss.

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

Run a validation step over the batch.

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

Run a test step over the batch.

on_train_epoch_end() → None

Aggregate and log training metrics and plots at epoch end.

on_train_epoch_start() → None

Reset accumulators at the start of a training epoch.

on_validation_epoch_start() → None

Reset accumulators at the start of a validation epoch.

on_validation_epoch_end() → None

Aggregate and log validation metrics and plots at epoch end.

configure_optimizers() → OptimizerLRSchedulerConfig

Build and return the optimizer and learning-rate scheduler.

The ReduceLROnPlateau scheduler is built and registered, but this task uses manual optimization and never calls lr_schedulers().step(), so Lightning does not step it and the learning rate stays at its configured value.