MolecularDiffusion.modules.models.syncogen.diffusion.loss.base

Base loss class for diffusion models.

Attributes

LossMode

Classes

LossBase	Base class for all loss functions with built-in coefficient and optional time weighting.
LossList	Aggregates multiple losses and returns dicts with individual values + totals.

Module Contents

class MolecularDiffusion.modules.models.syncogen.diffusion.loss.base.LossBase(mode: LossMode = 'both', coef: float = 1.0, time_weighted: bool = False, square_time_weight: bool = False, name: str = None, t_threshold: float = None)

Bases: torch.nn.Module

Base class for all loss functions with built-in coefficient and optional time weighting.

apply_coef(loss: torch.Tensor) → torch.Tensor

apply_t_threshold(loss: torch.Tensor, t: torch.Tensor | None) → torch.Tensor

Optionally mask loss to timesteps t <= t_threshold.

apply_time_weight(loss: torch.Tensor, t: torch.Tensor | None) → torch.Tensor

Optionally apply time weighting. Loss shapes are broadcast-compatible. NLLLoss should not call this (it has its own sigma weighting).

abstractmethod forward(*args, **kwargs)

coef = 1.0

mode = 'both'

property name: str

Return loss name. Defaults to class name if not explicitly set.

square_time_weight = False

t_threshold = None

time_weighted = False

class MolecularDiffusion.modules.models.syncogen.diffusion.loss.base.LossList(losses: Sequence[LossBase] = ())

Aggregates multiple losses and returns dicts with individual values + totals.

compute_coords(coords_pred, coords_gt, t) → dict[str, torch.Tensor]

compute_graph(log_p_X, log_p_E, node_mask, sigma_factor) → dict[str, torch.Tensor]

losses = []

MolecularDiffusion.modules.models.syncogen.diffusion.loss.base.LossMode