:py:mod:`anomalib.models.stfpm.loss`
====================================

.. py:module:: anomalib.models.stfpm.loss

.. autoapi-nested-parse::

   Loss function for the STFPM Model Implementation.



Module Contents
---------------

Classes
~~~~~~~

.. autoapisummary::

   anomalib.models.stfpm.loss.STFPMLoss




.. py:class:: STFPMLoss

   Bases: :py:obj:`torch.nn.Module`

   Feature Pyramid Loss This class implmenents the feature pyramid loss function proposed in STFPM paper.

   .. rubric:: Example

   >>> from anomalib.models.components.feature_extractors.feature_extractor import FeatureExtractor
   >>> from anomalib.models.stfpm.loss import STFPMLoss
   >>> from torchvision.models import resnet18

   >>> layers = ['layer1', 'layer2', 'layer3']
   >>> teacher_model = FeatureExtractor(model=resnet18(pretrained=True), layers=layers)
   >>> student_model = FeatureExtractor(model=resnet18(pretrained=False), layers=layers)
   >>> loss = Loss()

   >>> inp = torch.rand((4, 3, 256, 256))
   >>> teacher_features = teacher_model(inp)
   >>> student_features = student_model(inp)
   >>> loss(student_features, teacher_features)
       tensor(51.2015, grad_fn=<SumBackward0>)

   .. py:method:: compute_layer_loss(self, teacher_feats: torch.Tensor, student_feats: torch.Tensor) -> torch.Tensor

      Compute layer loss based on Equation (1) in Section 3.2 of the paper.

      :param teacher_feats: Teacher features
      :type teacher_feats: Tensor
      :param student_feats: Student features
      :type student_feats: Tensor

      :returns: L2 distance between teacher and student features.


   .. py:method:: forward(self, teacher_features: Dict[str, torch.Tensor], student_features: Dict[str, torch.Tensor]) -> torch.Tensor

      Compute the overall loss via the weighted average of the layer losses computed by the cosine similarity.

      :param teacher_features: Teacher features
      :type teacher_features: Dict[str, Tensor]
      :param student_features: Student features
      :type student_features: Dict[str, Tensor]

      :returns: Total loss, which is the weighted average of the layer losses.