torchcell.trainers.int_hetero_cell.DiffusionRegressionTask

class DiffusionRegressionTask(model: Module, cell_graph: Tensor, optimizer_config: dict[str, Any], lr_scheduler_config: dict[str, Any] | None, 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[Any, int] | None = None, device: str = 'cuda', inverse_transform: Module | None = None, execution_mode: str = 'training')

Bases: RegressionTask

RegressionTask for a diffusion decoder (006 hetero_cell_bipartite_dango_diff_gi).

loss_func must be a DiffusionLoss (or None for evaluation only): any other loss would train on the placeholders. Two differences, both from the model’s contract: in training mode GeneInteractionDiff returns all-zero placeholders instead of predictions (its diffusion loss trains the decoder from the targets and z_p), so train-stage predictions are not scored; and validation and test score the MSE of the sampled predictions, logged as {stage}/inference_mse, instead of calling the diffusion loss, which ignores the predictions. train/avg_diffusion_loss and val/avg_inference_mse are the epoch means of those per-batch losses.

on_train_epoch_end() → None

Log the epoch’s mean diffusion loss, then end the epoch as the base task.

on_validation_epoch_end() → None

Log the epoch’s mean inference MSE, then end the epoch as the base task.