| |
| from typing import Optional |
|
|
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch import Tensor |
|
|
| from mmdet.registry import MODELS |
| from .utils import weighted_loss |
|
|
|
|
| @weighted_loss |
| def knowledge_distillation_kl_div_loss(pred: Tensor, |
| soft_label: Tensor, |
| T: int, |
| detach_target: bool = True) -> Tensor: |
| r"""Loss function for knowledge distilling using KL divergence. |
| |
| Args: |
| pred (Tensor): Predicted logits with shape (N, n + 1). |
| soft_label (Tensor): Target logits with shape (N, N + 1). |
| T (int): Temperature for distillation. |
| detach_target (bool): Remove soft_label from automatic differentiation |
| |
| Returns: |
| Tensor: Loss tensor with shape (N,). |
| """ |
| assert pred.size() == soft_label.size() |
| target = F.softmax(soft_label / T, dim=1) |
| if detach_target: |
| target = target.detach() |
|
|
| kd_loss = F.kl_div( |
| F.log_softmax(pred / T, dim=1), target, reduction='none').mean(1) * ( |
| T * T) |
|
|
| return kd_loss |
|
|
|
|
| @MODELS.register_module() |
| class KnowledgeDistillationKLDivLoss(nn.Module): |
| """Loss function for knowledge distilling using KL divergence. |
| |
| Args: |
| reduction (str): Options are `'none'`, `'mean'` and `'sum'`. |
| loss_weight (float): Loss weight of current loss. |
| T (int): Temperature for distillation. |
| """ |
|
|
| def __init__(self, |
| reduction: str = 'mean', |
| loss_weight: float = 1.0, |
| T: int = 10) -> None: |
| super().__init__() |
| assert T >= 1 |
| self.reduction = reduction |
| self.loss_weight = loss_weight |
| self.T = T |
|
|
| def forward(self, |
| pred: Tensor, |
| soft_label: Tensor, |
| weight: Optional[Tensor] = None, |
| avg_factor: Optional[int] = None, |
| reduction_override: Optional[str] = None) -> Tensor: |
| """Forward function. |
| |
| Args: |
| pred (Tensor): Predicted logits with shape (N, n + 1). |
| soft_label (Tensor): Target logits with shape (N, N + 1). |
| weight (Tensor, optional): The weight of loss for each |
| prediction. Defaults to None. |
| avg_factor (int, optional): Average factor that is used to average |
| the loss. Defaults to None. |
| reduction_override (str, optional): The reduction method used to |
| override the original reduction method of the loss. |
| Defaults to None. |
| |
| Returns: |
| Tensor: Loss tensor. |
| """ |
| assert reduction_override in (None, 'none', 'mean', 'sum') |
|
|
| reduction = ( |
| reduction_override if reduction_override else self.reduction) |
|
|
| loss_kd = self.loss_weight * knowledge_distillation_kl_div_loss( |
| pred, |
| soft_label, |
| weight, |
| reduction=reduction, |
| avg_factor=avg_factor, |
| T=self.T) |
|
|
| return loss_kd |
|
|
