torchcell.trainers.dcell_regression_slim.DCellRegressionSlimTask

class DCellRegressionSlimTask(models: dict[str, Module], target: str, boxplot_every_n_epochs: int = 10, learning_rate: float = 0.001, weight_decay: float = 1e-05, batch_size: int | None = None, alpha: float = 0.3, *, aux_reduction: Literal['sum', 'mean'], **kwargs: Any)

Bases: LightningModule

LightningModule for training models on graph-based regression datasets.

setup(stage: str | None = None) → None

Move all submodels to the active device at the start of each stage.

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

Run the batch through the DCell subsystems and linear head.

on_train_start() → None

Log the total model parameter count when training starts.

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

Run a manual training step and log loss plus subsystem and root metrics.

on_train_epoch_end() → None

Log and reset the subsystem and root training metrics at epoch end.

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

Run a validation step and update subsystem and root metrics.

on_validation_epoch_end() → None

Log and reset the subsystem and root validation metrics.

on_train_epoch_start() → None

Log the checkpoint the previous epoch saved, if it is a new best.

on_train_end() → None

Log the checkpoint the final epoch saved, if it is a new best.

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

Run a test step and update subsystem and root metrics.

on_test_epoch_end() → None

Log and reset the subsystem and root test metrics at epoch end.

configure_optimizers() → Optimizer

Return an Adam optimizer over the DCell and linear-head parameters.