autogluon.multimodal 1.2.1b20250303__py3-none-any.whl → 1.2.1b20250304__py3-none-any.whl

autogluon/multimodal/optim/__init__.py

@@ -0,0 +1,17 @@
+from .lit_distiller import DistillerLitModule
+from .lit_matcher import MatcherLitModule
+from .lit_mmdet import MMDetLitModule
+from .lit_module import LitModule
+from .lit_ner import NerLitModule
+from .lit_semantic_seg import SemanticSegmentationLitModule
+from .losses import get_aug_loss_func, get_loss_func, get_matcher_loss_func, get_matcher_miner_func
+from .metrics import (
+    CustomHitRate,
+    compute_ranking_score,
+    compute_score,
+    get_minmax_mode,
+    get_stopping_threshold,
+    get_torchmetric,
+    infer_metrics,
+)
+from .utils import get_norm_layer_param_names, get_peft_param_names

autogluon/multimodal/{optimization → optim}/lit_distiller.py

@@ -13,7 +13,8 @@ from torchmetrics.aggregation import BaseAggregator
 
 from ..constants import FEATURES, LOGITS, WEIGHT
 from ..models.utils import run_model
-from .utils import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler, get_optimizer
+from .lr import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler
+from .utils import get_optimizer
 
 logger = logging.getLogger(__name__)
 

autogluon/multimodal/{optimization → optim}/lit_matcher.py

@@ -13,16 +13,10 @@ from torchmetrics.aggregation import BaseAggregator
 from ..constants import FEATURES, LOGIT_SCALE, PROBABILITY, QUERY, RESPONSE
 from ..models.utils import run_model
 from ..utils.matcher import compute_matching_probability
-from .losses import MultiNegativesSoftmaxLoss
-from .utils import (
-    CustomHitRate,
-    apply_layerwise_lr_decay,
-    apply_single_lr,
-    apply_two_stages_lr,
-    generate_metric_learning_labels,
-    get_lr_scheduler,
-    get_optimizer,
-)
+from .losses import MultiNegativesSoftmaxLoss, generate_metric_learning_labels
+from .lr import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler
+from .metrics import CustomHitRate
+from .utils import get_optimizer
 
 logger = logging.getLogger(__name__)
 

autogluon/multimodal/{optimization → optim}/lit_mmdet.py

@@ -2,21 +2,13 @@ import logging
 from typing import Callable, Optional, Union
 
 import lightning.pytorch as pl
-import torch
 import torchmetrics
 from lightning.pytorch.utilities import grad_norm
-from torch.nn.modules.loss import _Loss
 from torchmetrics.aggregation import BaseAggregator
 
 from ..constants import BBOX, IMAGE, LABEL
-from .utils import (
-    apply_layerwise_lr_decay,
-    apply_single_lr,
-    apply_two_stages_lr,
-    get_lr_scheduler,
-    get_optimizer,
-    remove_parameters_without_grad,
-)
+from .lr import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler
+from .utils import get_optimizer, remove_parameters_without_grad
 
 try:
     import mmdet

autogluon/multimodal/{optimization → optim}/lit_module.py

@@ -11,11 +11,25 @@ from torch import nn
 from torch.nn.modules.loss import _Loss
 from torchmetrics.aggregation import BaseAggregator
 
-from ..constants import LM_TARGET, LOGITS, T_FEW, TEMPLATE_LOGITS, WEIGHT
+from ..constants import (
+    AUG_LOGITS,
+    LM_TARGET,
+    LOGITS,
+    MULTIMODAL_FEATURES,
+    MULTIMODAL_FEATURES_POST_AUG,
+    MULTIMODAL_FEATURES_PRE_AUG,
+    ORI_LOGITS,
+    T_FEW,
+    TEMPLATE_LOGITS,
+    VAE_MEAN,
+    VAE_VAR,
+    WEIGHT,
+)
 from ..data.mixup import MixupModule, multimodel_mixup
 from ..models.utils import run_model
-from .semantic_seg_metrics import COD, Balanced_Error_Rate
-from .utils import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler, get_optimizer
+from .lr import apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, get_lr_scheduler
+from .metrics import Coverage
+from .utils import get_optimizer
 
 logger = logging.getLogger(__name__)
 
@@ -44,13 +58,24 @@ class LitModule(pl.LightningModule):
         validation_metric_name: Optional[str] = None,
         custom_metric_func: Callable = None,
         test_metric: Optional[torchmetrics.Metric] = None,
-        efficient_finetune: Optional[str] = None,
+        peft: Optional[str] = None,
         trainable_param_names: Optional[List] = None,
         mixup_fn: Optional[MixupModule] = None,
         mixup_off_epoch: Optional[int] = 0,
         model_postprocess_fn: Callable = None,
         skip_final_val: Optional[bool] = False,
         track_grad_norm: Optional[Union[int, str]] = -1,
+        cross_modal_align: Optional[str] = None,
+        cross_modal_align_weight: Optional[float] = 0,
+        automatic_optimization: Optional[bool] = True,
+        accumulate_grad_batches: Optional[int] = None,
+        gradient_clip_val: Optional[float] = None,
+        gradient_clip_algorithm: Optional[str] = None,
+        use_aug_optim: Optional[bool] = False,
+        aug_loss_func: Optional[_Loss] = None,
+        aug_lr: Optional[float] = None,
+        aug_weight_decay: Optional[float] = None,
+        aug_optim_type: Optional[str] = None,
     ):
         """
         Parameters
@@ -104,7 +129,7 @@ class LitModule(pl.LightningModule):
             Refer to https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/aggregation.py
         test_metric
             A torchmetrics module used in the test stage, e.g., torchmetrics.Accuracy().
-        efficient_finetune
+        peft
             Whether to use efficient finetuning strategies. This will be helpful for fast finetuning of large backbones.
             We support options such as:
 
@@ -128,6 +153,8 @@ class LitModule(pl.LightningModule):
                 "model_postprocess_fn",
                 "mixup_fn",
                 "trainable_param_names",
+                "custom_metric_func",
+                "aug_loss_func",
             ]
         )
         self.model = model
@@ -144,7 +171,12 @@ class LitModule(pl.LightningModule):
         self.model_postprocess_fn = model_postprocess_fn
         self.trainable_param_names = trainable_param_names if trainable_param_names else []
         self.skip_final_val = skip_final_val
-        self.track_grad_norm = track_grad_norm
+        self.automatic_optimization = automatic_optimization
+        self.aug_loss_func = aug_loss_func
+        if self.hparams.cross_modal_align:
+            assert self.hparams.cross_modal_align_weight > 0
+            logger.debug(f"Cross modal alignment mode: {self.hparams.cross_modal_align}")
+            logger.debug(f"Cross modal alignment loss weight: {self.hparams.cross_modal_align_weight}")
 
     def _compute_template_loss(
         self,
@@ -179,19 +211,41 @@ class LitModule(pl.LightningModule):
 
         return lm_loss + mc_loss * self.model.mc_loss + unlikely_loss * self.model.unlikely_loss
 
+    def _compute_cross_modal_align_loss(self, multimodal_features):
+        if self.hparams.cross_modal_align == "positive_only":
+            kl_loss = nn.KLDivLoss(reduction="batchmean", log_target=True)
+            loss = 0
+            num = 0
+            for i in range(len(multimodal_features)):
+                # input should be a distribution in the log space
+                a = F.log_softmax(multimodal_features[i], dim=1)
+                # kl divergence is not symmetric, so need to compute both (i, j) and (j, i)
+                for j in range(len(multimodal_features)):
+                    if i == j:
+                        continue
+                    # input should be a distribution in the log space
+                    b = F.log_softmax(multimodal_features[j], dim=1)
+                    loss += kl_loss(a, b)
+                    num += 1
+            return self.hparams.cross_modal_align_weight * loss / num
+        else:
+            raise ValueError(f"Unsupported cross modal alignment loss: {self.hparams.cross_modal_align}.")
+
     def _compute_loss(
         self,
         output: Dict,
         label: torch.Tensor,
     ):
         loss = 0
-        for _, per_output in output.items():
+        for per_prefix, per_output in output.items():
             weight = per_output[WEIGHT] if WEIGHT in per_output else 1
             if (
                 TEMPLATE_LOGITS in per_output and self.model.prefix == T_FEW
             ):  # Do only add template loss if T-Few. #TODO Add compatibility to Fusion models.
                 loss += self._compute_template_loss(per_output, label) * weight
             else:
+                if self.training and self.hparams.use_aug_optim and per_prefix.startswith("fusion"):
+                    label = label.tile((2,))
                 loss += (
                     self.loss_func(
                         input=per_output[LOGITS].squeeze(dim=1),
@@ -199,6 +253,22 @@ class LitModule(pl.LightningModule):
                     )
                     * weight
                 )
+
+        if self.hparams.cross_modal_align:
+            loss += self._compute_cross_modal_align_loss(
+                multimodal_features=output[self.model.prefix][MULTIMODAL_FEATURES]
+            )
+
+        if self.training and self.hparams.use_aug_optim:
+            loss += self.aug_loss_func(
+                pre_aug=output[self.model.prefix][MULTIMODAL_FEATURES_PRE_AUG],
+                post_aug=output[self.model.prefix][MULTIMODAL_FEATURES_POST_AUG],
+                vae_mean=output[self.model.prefix][VAE_MEAN],
+                vae_var=output[self.model.prefix][VAE_VAR],
+                ori_logits=output[self.model.prefix][ORI_LOGITS],
+                aug_logits=output[self.model.prefix][AUG_LOGITS],
+            )
+
         return loss
 
     def _compute_metric_score(
@@ -214,6 +284,7 @@ class LitModule(pl.LightningModule):
                 torchmetrics.classification.BinaryAUROC,
                 torchmetrics.classification.BinaryAveragePrecision,
                 torchmetrics.classification.BinaryF1Score,
+                Coverage,
             ),
         ):
             prob = F.softmax(logits.float(), dim=1)
@@ -255,6 +326,27 @@ class LitModule(pl.LightningModule):
         Average loss of the mini-batch data.
         """
         output, loss = self._shared_step(batch)
+
+        if not self.automatic_optimization:
+            if self.hparams.use_aug_optim:
+                optimizer, aug_optimizer = self.optimizers()
+            else:
+                optimizer = self.optimizers()
+                aug_optimizer = None
+
+            lr_scheduler = self.lr_schedulers()
+            loss = loss / self.hparams.accumulate_grad_batches
+            self.manual_backward(loss)
+
+            if (batch_idx + 1) % self.hparams.accumulate_grad_batches == 0 or self.trainer.is_last_batch:
+                optimizer.step()
+                optimizer.zero_grad()
+                lr_scheduler.step()
+
+                if aug_optimizer is not None:
+                    aug_optimizer.step()
+                    aug_optimizer.zero_grad()
+
         self.log("train_loss", loss)
         return loss
 
@@ -345,6 +437,9 @@ class LitModule(pl.LightningModule):
             model=self.model,
             lr=self.hparams.lr,
             weight_decay=self.hparams.weight_decay,
+            exclude_keys=[
+                "augmenter"
+            ],  # exclude augmenter parameters from the optimizer as they would use an independent optimizer
         )
         if self.hparams.lr_choice == "two_stages":
             logger.debug("applying 2-stage learning rate...")
@@ -357,14 +452,14 @@ class LitModule(pl.LightningModule):
             logger.debug("applying layerwise learning rate decay...")
             grouped_parameters = apply_layerwise_lr_decay(
                 lr_decay=self.hparams.lr_decay,
-                efficient_finetune=self.hparams.efficient_finetune,
+                peft=self.hparams.peft,
                 trainable_param_names=self.trainable_param_names,
                 **kwargs,
             )
         else:
             logger.debug("applying single learning rate...")
             grouped_parameters = apply_single_lr(
-                efficient_finetune=self.hparams.efficient_finetune,
+                peft=self.hparams.peft,
                 trainable_param_names=self.trainable_param_names,
                 **kwargs,
             )
@@ -381,16 +476,21 @@ class LitModule(pl.LightningModule):
             if isinstance(self.trainer.strategy, DeepSpeedStrategy):
                 max_steps = 1
             else:
+                accumulate_grad_batches = (
+                    self.trainer.accumulate_grad_batches
+                    if self.automatic_optimization
+                    else self.hparams.accumulate_grad_batches
+                )
                 max_steps = (
                     len(self.trainer.datamodule.train_dataloader())
                     * self.trainer.max_epochs
-                    // self.trainer.accumulate_grad_batches
+                    // accumulate_grad_batches
                 )
                 logger.debug(
                     f"len(trainer.datamodule.train_dataloader()): {len(self.trainer.datamodule.train_dataloader())}"
                 )
                 logger.debug(f"trainer.max_epochs: {self.trainer.max_epochs}")
-                logger.debug(f"trainer.accumulate_grad_batches: {self.trainer.accumulate_grad_batches}")
+                logger.debug(f"accumulate_grad_batches: {accumulate_grad_batches}")
         else:
             max_steps = self.trainer.max_steps
 
@@ -411,10 +511,35 @@ class LitModule(pl.LightningModule):
         )
 
         sched = {"scheduler": scheduler, "interval": "step"}
+        ret_optimizers = [optimizer]
+        ret_schedulers = [sched]
+        if self.hparams.use_aug_optim:
+            logger.debug("initializing augment optimizer")
+            # augmenter's optimizer
+            aug_grouped_parameters = apply_single_lr(
+                model=self.model.augmenter,
+                lr=self.hparams.aug_lr,
+                weight_decay=self.hparams.aug_weight_decay,
+            )
+            aug_optimizer = get_optimizer(
+                optim_type=self.hparams.aug_optim_type,
+                optimizer_grouped_parameters=aug_grouped_parameters,
+                lr=self.hparams.aug_lr,
+                weight_decay=self.hparams.aug_weight_decay,
+            )
+            ret_optimizers.append(aug_optimizer)
+
         logger.debug("done configuring optimizer and scheduler")
-        return [optimizer], [sched]
+        return ret_optimizers, ret_schedulers
 
     def on_before_optimizer_step(self, optimizer):
         # If using mixed precision, the gradients are already unscaled here
-        if self.track_grad_norm != -1:
-            self.log_dict(grad_norm(self, norm_type=self.track_grad_norm))
+        # TODO: apply gradient clip only to the target optimizer
+        if not self.automatic_optimization and self.hparams.gradient_clip_val > 0:
+            self.clip_gradients(
+                optimizer,
+                gradient_clip_val=self.hparams.gradient_clip_val,
+                gradient_clip_algorithm=self.hparams.gradient_clip_algorithm,
+            )
+        if self.hparams.track_grad_norm != -1:
+            self.log_dict(grad_norm(self, norm_type=self.hparams.track_grad_norm))

autogluon/multimodal/{optimization → optim}/lit_ner.py

@@ -37,7 +37,7 @@ class NerLitModule(LitModule):
         validation_metric_name: Optional[str] = None,
         custom_metric_func: Callable = None,
         test_metric: Optional[torchmetrics.Metric] = None,
-        efficient_finetune: Optional[str] = None,
+        peft: Optional[str] = None,
         trainable_param_names: Optional[List] = None,
         mixup_fn: Optional[MixupModule] = None,
         mixup_off_epoch: Optional[int] = 0,
@@ -97,7 +97,7 @@ class NerLitModule(LitModule):
             Refer to https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/aggregation.py
         test_metric
             A torchmetrics module used in the test stage, e.g., torchmetrics.Accuracy().
-        efficient_finetune
+        peft
             Whether to use efficient finetuning strategies. This will be helpful for fast finetuning of large backbones.
             We support options such as:
 
@@ -127,7 +127,7 @@ class NerLitModule(LitModule):
             validation_metric_name=validation_metric_name,
             custom_metric_func=custom_metric_func,
             test_metric=test_metric,
-            efficient_finetune=efficient_finetune,
+            peft=peft,
             trainable_param_names=trainable_param_names,
             mixup_fn=mixup_fn,
             mixup_off_epoch=mixup_off_epoch,

autogluon/multimodal/{optimization → optim}/lit_semantic_seg.py

@@ -8,7 +8,7 @@ from transformers.models.mask2former.modeling_mask2former import Mask2FormerLoss
 from ..constants import CLASS_LOGITS, LOGITS, MOE_LOSS, SEMANTIC_MASK, WEIGHT
 from ..models.utils import run_model
 from .lit_module import LitModule
-from .semantic_seg_metrics import Multiclass_IoU
+from .metrics.semantic_seg_metrics import Multiclass_IoU
 
 logger = logging.getLogger(__name__)
 

autogluon/multimodal/optim/losses/__init__.py

@@ -0,0 +1,14 @@
+from .softmax_losses import MultiNegativesSoftmaxLoss, SoftTargetCrossEntropy
+from .focal_loss import FocalLoss
+from .lemda_loss import LemdaLoss
+from .rkd_loss import RKDLoss
+from .bce_loss import BBCEWithLogitLoss
+from .structure_loss import StructureLoss
+from .utils import (
+    generate_metric_learning_labels,
+    get_aug_loss_func,
+    get_loss_func,
+    get_matcher_loss_func,
+    get_matcher_miner_func,
+    get_metric_learning_distance_func,
+)

autogluon/multimodal/optim/losses/bce_loss.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+
+
+class BBCEWithLogitLoss(nn.Module):
+    """
+    Balanced BCEWithLogitLoss based on https://github.com/NiFangBaAGe/Explicit-Visual-Prompt/blob/latest_branch/models/segformer.py
+    """
+
+    def __init__(self):
+        super(BBCEWithLogitLoss, self).__init__()
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor):
+        if input.dim() == 3:
+            input = input.unsqueeze(1)
+        eps = 1e-10
+        count_pos = torch.sum(target) + eps
+        count_neg = torch.sum(1.0 - target)
+        ratio = count_neg / count_pos
+        w_neg = count_pos / (count_pos + count_neg)
+
+        bce1 = nn.BCEWithLogitsLoss(pos_weight=ratio)
+        loss = w_neg * bce1(input, target)
+
+        return loss

autogluon/multimodal/optim/losses/focal_loss.py

@@ -0,0 +1,81 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FocalLoss(nn.Module):
+    """
+    Focal loss based on https://github.com/AdeelH/pytorch-multi-class-focal-loss/blob/master/focal_loss.py
+
+    References:
+        [1] https://arxiv.org/abs/1708.02002
+    """
+
+    def __init__(
+        self,
+        alpha: Optional[torch.Tensor] = None,
+        gamma: Optional[float] = 2.0,
+        reduction: Optional[str] = "mean",
+        eps: Optional[float] = 1e-6,
+    ):
+        """
+
+        Parameters
+        ----------
+        alpha
+            weighting factor for each class. Should be of shape (num_classes)
+        gamma
+            the focal parameter for calculating weights on easy/hard samples
+        reduction
+            the reduction to apply to the final loss output. Default: "mean". Options:
+                "mean", "sum"
+        eps
+            epsilon for numerical stability
+        """
+        super(FocalLoss, self).__init__()
+
+        self.gamma = gamma
+        self.reduction = reduction
+        self.eps = eps
+        if alpha is not None:
+            if isinstance(alpha, str):  # handles Ray Tune HPO sampled hyperparameter
+                try:
+                    numbers = alpha.strip("()").split(",")
+                    alpha = [float(num) for num in numbers]
+                except:
+                    raise ValueError(f"{type(alpha)} {alpha} is not in a supported format.")
+            alpha = torch.tensor(alpha)
+        self.nll_loss = nn.NLLLoss(weight=alpha, reduction="none")
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor):
+        if not torch.is_tensor(input):
+            raise TypeError("input type is not a torch.Tensor. Got {}".format(type(input)))
+        if input.ndim > 2:
+            # (N, C, d1, d2, ..., dK) --> (N * d1 * ... * dK, C)
+            num_class = input.shape[1]
+            input = input.permute(0, *range(2, input.ndim), 1).reshape(-1, num_class)
+            # (N, d1, d2, ..., dK) --> (N * d1 * ... * dK,)
+            target = target.view(-1)
+
+        pt = F.softmax(input, dim=-1)
+
+        # -alpha_t * log(pt) term
+        log_p = torch.log_softmax(input, dim=-1)
+        ce = self.nll_loss(log_p, target)
+
+        # (1 - pt)^gamma term
+        all_rows = torch.arange(input.shape[0])
+        pt = pt[all_rows, target]
+        focal_term = (1 - pt) ** self.gamma
+
+        loss = focal_term * ce
+
+        if self.reduction == "mean":
+            loss = loss.mean()
+
+        elif self.reduction == "sum":
+            loss = loss.sum()
+
+        return loss

autogluon/multimodal/optim/losses/lemda_loss.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...constants import BINARY, REGRESSION
+
+
+class LemdaLoss(nn.Module):
+    def __init__(self, mse_weight, kld_weight, consist_weight, consist_threshold, problem_type):
+        super().__init__()
+        self.mse_loss = nn.MSELoss(reduction="mean")
+        self.mse_weight = mse_weight
+        self.kld_weight = kld_weight
+        self.consist_weight = consist_weight
+        self.consist_threshold = consist_threshold
+        self.problem_type = problem_type
+
+    def consist_loss(self, p_logits, q_logits):
+        p = F.softmax(p_logits, dim=1)
+        logp = F.log_softmax(p_logits, dim=1)
+        logq = F.log_softmax(q_logits, dim=1)
+        loss = torch.sum(p * (logp - logq), dim=-1)
+        q = F.softmax(q_logits, dim=1)
+        q_largest = torch.max(q, dim=1)[0]
+        loss_mask = torch.gt(q_largest, self.consist_threshold).float()
+        loss = loss * loss_mask
+        return torch.mean(loss)
+
+    def forward(self, pre_aug, post_aug, vae_mean, vae_var, ori_logits, aug_logits):
+        mse_loss = self.mse_loss(pre_aug, post_aug) * self.mse_weight
+        # see Appendix B from VAE paper:    https://arxiv.org/abs/1312.6114
+        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
+        kld_loss = -0.5 * torch.mean(1 + vae_var - vae_mean.pow(2) - vae_var.exp()) * self.kld_weight
+        if self.problem_type in [REGRESSION, BINARY]:
+            consist_loss = self.mse_loss(ori_logits, aug_logits) * self.consist_weight
+        else:
+            consist_loss = self.consist_loss(ori_logits, aug_logits) * self.consist_weight
+
+        return mse_loss + kld_loss + consist_loss

autogluon/multimodal/optim/losses/rkd_loss.py

@@ -0,0 +1,103 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RKDLoss(nn.Module):
+    """
+    Compute RKD Distance Loss.
+    Paper Refer to: Relational Knowledge Disitllation, CVPR2019. https://arxiv.org/abs/1904.05068
+    Code Refer to: https://github.com/HobbitLong/RepDistiller/blob/master/distiller_zoo/RKD.py
+    and https://github.com/lenscloth/RKD/blob/master/metric/loss.py
+    """
+
+    def __init__(self, distance_loss_weight: Optional[float] = 25.0, angle_loss_weight: Optional[float] = 50.0):
+        """
+        Parameters
+        ----------
+        distance_loss_weight
+            Weight of RKD distance loss
+        angle_loss_weight
+            Weight of RKD angle loss
+        Returns
+        -------
+        """
+        super(RKDLoss, self).__init__()
+        self.distance_loss_weight = distance_loss_weight
+        self.angle_loss_weight = angle_loss_weight
+
+    def forward(self, feature_student: Optional[torch.Tensor], feature_teacher: Optional[torch.Tensor]):
+        """
+        Parameters
+        ----------
+        feature_student
+            Output feature of student model, shape: (N, D)
+        feature_teacher
+            Output feature of teacher model, shape: (N, D)
+        Returns
+        -------
+        The RKD Loss between teacher and student
+        """
+        # RKD loss
+        if self.distance_loss_weight > 0:
+            with torch.no_grad():
+                t_dist = self.pdist(feature_teacher, squared=False)
+                mean_td = t_dist[t_dist > 0].mean()
+                t_dist = t_dist / mean_td
+
+            s_dist = self.pdist(feature_student, squared=False)
+            mean_d = s_dist[s_dist > 0].mean()
+            s_dist = s_dist / mean_d
+
+            loss_distance = F.smooth_l1_loss(s_dist, t_dist)
+
+        # RKD Angle loss
+        if self.angle_loss_weight > 0:
+            with torch.no_grad():
+                td = feature_teacher.unsqueeze(0) - feature_teacher.unsqueeze(1)
+                norm_td = F.normalize(td, p=2, dim=2)
+                t_angle = torch.bmm(norm_td, norm_td.transpose(1, 2)).view(-1)
+
+            sd = feature_student.unsqueeze(0) - feature_student.unsqueeze(1)
+            norm_sd = F.normalize(sd, p=2, dim=2)
+            s_angle = torch.bmm(norm_sd, norm_sd.transpose(1, 2)).view(-1)
+
+            loss_angle = F.smooth_l1_loss(s_angle, t_angle)
+
+        loss = ((self.distance_loss_weight * loss_distance) if self.distance_loss_weight > 0 else 0) + (
+            (self.angle_loss_weight * loss_angle) if self.angle_loss_weight > 0 else 0
+        )
+
+        return loss
+
+    @staticmethod
+    def pdist(embeddings: Optional[torch.Tensor], squared: Optional[bool] = False, eps: Optional[float] = 1e-12):
+        """
+        Compute pairwise Euclidean distances between embeddings in n-dimensional space.
+
+        Parameters
+        ----------
+        embeddings
+            The embeddings to compute pairwise distance between. Shape: (N,D)
+        squared
+            If the result is square of Euclidean distance.
+        eps
+            Min value of each entry.
+
+        Returns
+        -------
+        Pairwise Euclidean distances. Shape: (N,N)
+        """
+        e_square = embeddings.pow(2).sum(dim=1)
+        prod = embeddings @ embeddings.t()
+        res = (e_square.unsqueeze(1) + e_square.unsqueeze(0) - 2 * prod).clamp(min=eps)
+
+        if not squared:
+            res = res.sqrt()
+
+        res = res.clone()
+        res[range(len(embeddings)), range(len(embeddings))] = 0
+
+        return res