Class Loss

flowmap/loss/loss.py:24–58

class Loss(nn.Module, ABC, Generic[T]):
    cfg: T

    def __init__(self, cfg: T) -> None:
        super().__init__()
        self.cfg = cfg

    def forward(
        self,
        batch: Batch,
        flows: Flows,
        tracks: list[Tracks] | None,
        model_output: ModelOutput,
        global_step: int,
    ) -> Float[Tensor, ""]:
        # Before the loss is enabled, don't compute the loss.
        if global_step < self.cfg.enable_after:
            return torch.tensor(0, dtype=torch.float32, device=batch.videos.device)

        # Multiply the computed loss value by the weight.
        loss = self.compute_unweighted_loss(
            batch, flows, tracks, model_output, global_step
        )
        return self.cfg.weight * loss

    @abstractmethod
    def compute_unweighted_loss(
        self,
        batch: Batch,
        flows: Flows,
        tracks: list[Tracks] | None,
        model_output: ModelOutput,
        global_step: int,
    ) -> Float[Tensor, ""]:
        pass