torchcell.trainers.dcell_regression.DCellRegressionTask

class DCellRegressionTask(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, train_wt_diff: bool = True, *, aux_reduction: Literal['sum', 'mean'], **kwargs: object)

Bases: LightningModule

LightningModule for training models on graph-based regression datasets.

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

Move submodels to device and init prediction/true-value buffers.

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

Run the DCell subsystem and linear heads, returning per-node outputs.

on_train_start() → None

Log the total parameter count at the start of training.

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

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

on_train_epoch_end() → None

Log and reset accumulated training metrics.

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

Run a validation step, logging loss/metrics and storing predictions.

on_validation_epoch_end() → None

Log val metrics and, on plotting epochs, the root-prediction box plot.

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 log loss, regression metrics, and correlations.

on_test_epoch_end() → None

Log and reset accumulated test metrics.

configure_optimizers() → Optimizer

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