Toggle Light / Dark / Auto color theme

Toggle table of contents sidebar

Source code for anomalib.models.fastflow.loss

"""Loss function for the FastFlow Model Implementation."""

# Copyright (C) 2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from typing import List

import torch
from torch import Tensor, nn


class FastflowLoss(nn.Module):
    """FastFlow Loss."""

    def forward(self, hidden_variables: List[Tensor], jacobians: List[Tensor]) -> Tensor:
        """Calculate the Fastflow loss.

        Args:
            hidden_variables (List[Tensor]): Hidden variables from the fastflow model. f: X -> Z
            jacobians (List[Tensor]): Log of the jacobian determinants from the fastflow model.

        Returns:
            Tensor: Fastflow loss computed based on the hidden variables and the log of the Jacobians.
        """
        loss = torch.tensor(0.0, device=hidden_variables[0].device)  # pylint: disable=not-callable
        for (hidden_variable, jacobian) in zip(hidden_variables, jacobians):
            loss += torch.mean(0.5 * torch.sum(hidden_variable**2, dim=(1, 2, 3)) - jacobian)
        return loss