torchcell.trainers.int_dcell.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)

Bases: LightningModule

Lightning module training a DCell model to regress gene interactions.

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

Run a forward pass through the model on the given batch.

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

Run a manual-optimization training step with gradient accumulation.

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

Run the shared step for the validation stage and return the loss.

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

Run the shared step for the test stage and return the loss.

on_train_epoch_end() → None

Log and reset training metrics and plot training samples periodically.

on_train_epoch_start() → None

Clear training sample buffers on epochs where samples are collected.

on_validation_epoch_start() → None

Clear validation sample buffers on epochs where samples are collected.

on_test_epoch_start() → None

Clear test sample buffers at the start of the test epoch.

on_validation_epoch_end() → None

Log and reset validation metrics and plot validation samples periodically.

on_test_epoch_end() → None

Log and reset test metrics and plot test samples.

configure_optimizers() → OptimizerLRSchedulerConfig

Set up optimizers for training.

For DCellModel, we need to pre-register parameters by running a forward pass.