torchcell.trainers.simple_linear_regression.SimpleLinearRegressionTask

class SimpleLinearRegressionTask(model: Module, target: str, boxplot_every_n_epochs: int = 10, 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, **kwargs: Any)

Bases: LightningModule

LightningModule for training models on graph-based regression datasets.

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

Move the model to the trainer’s device.

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

Run the model and return squeezed predictions.

on_train_start() → None

Log the total number of model parameters at training start.

training_step(batch: Any, 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: Any, batch_idx: int) → None

Run a validation step, logging loss, metrics, and correlations.

on_validation_epoch_end() → None

Log metrics and, on plotting epochs, render the 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: Any, batch_idx: int) → None

Run a test step, logging loss, metrics, and correlations.

on_test_epoch_end() → None

Log and reset accumulated test metrics.

configure_optimizers() → Optimizer

Return an Adam optimizer over the model parameters.