torchcell.trainers.neo_regression.RegressionTask

class RegressionTask(model: Module, target: str, learning_rate: float = 0.001, weight_decay: float = 1e-05, loss: str = 'mse', batch_size: int | None = None, train_epoch_size: int | None = None, clip_grad_norm: bool = False, clip_grad_norm_max_norm: float = 0.1, boxplot_every_n_epochs: int = 1, alpha: float = 0.0)

Bases: LightningModule

LightningModule for training models on graph-based regression datasets.

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

Move the model to the active device at the start of each stage.

forward(x: Tensor, batch: Tensor) → Tensor

Run inputs through the main encoder and top head to produce predictions.

on_train_start() → None

Log the total model parameter count when training starts.

training_step(batch: dict[str, Any], batch_idx: int) → Tensor

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

on_train_epoch_end() → None

Log and reset the accumulated training metrics at epoch end.

validation_step(batch: dict[str, Any], batch_idx: int) → None

Run a validation step, log metrics, and store predictions for plotting.

compute_prediction_stats(true_values: Tensor, predictions: Tensor, stage: str = 'val') → None

Bin predictions and log per-bin mean and std as a wandb table.

on_validation_epoch_end() → None

Log validation metrics, draw box plots, and log the best model artifact.

test_step(batch: dict[str, Any], batch_idx: int) → None

Run a test step, log metrics, and store predictions for plotting.

on_test_epoch_end() → None

Log test metrics and draw the binned-prediction box plot.

configure_optimizers() → Optimizer

Return an Adam optimizer over the model parameters.