Update main

Author: lizekai-richard

dd_ranking/aug/cutmix.py

@@ -5,15 +5,6 @@
 
 class Cutmix_Augmentation:
     def __init__(self, params: dict):
-        # self.transform = kornia.augmentation.RandomCutMixV2(
-        #     num_mix = params["times"],
-        #     cut_size = params["size"],
-        #     same_on_batch = params["same_on_batch"],
-        #     beta = params["beta"],
-        #     keepdim = params["keep_dim"],
-        #     p = params["prob"]
-        # )
-
         self.cutmix_p = params["cutmix_p"]
 
     def rand_bbox(self, size, lam):

dd_ranking/aug/mixup.py

@@ -5,13 +5,6 @@
 
 class Mixup_Augmentation:
     def __init__(self, params: dict):
-        # self.transform = kornia.augmentation.RandomMixUpV2(
-        #     lambda_val = params["lambda_range"],
-        #     same_on_batch = params["same_on_batch"],
-        #     keepdim = params["keepdim"],
-        #     p = params["prob"]
-        # )
-
         self.mixup_p = params["mixup_p"]
 
     def mixup(self, images):

dd_ranking/loss/kl.py

@@ -1,17 +1,16 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.nn import KLDivLoss
 
 
 class KLDivergenceLoss(nn.Module):
     def __init__(self, temperature=1.2):
         super(KLDivergenceLoss, self).__init__()
         self.temperature = temperature
-        self.kl = KLDivLoss(reduction='batchmean')
 
     def forward(self, stu_outputs, tea_outputs):
         stu_probs = F.log_softmax(stu_outputs / self.temperature, dim=1)
-        tea_probs = F.softmax(tea_outputs / self.temperature, dim=1)
-        loss = self.kl(stu_probs, tea_probs)
+        with torch.no_grad():
+            tea_probs = F.softmax(tea_outputs / self.temperature, dim=1)
+        loss = F.kl_div(stu_probs, tea_probs, reduction='batchmean')
         return loss

dd_ranking/loss/sce.py

@@ -4,12 +4,13 @@
 
 
 class SoftCrossEntropyLoss(nn.Module):
-    def __init__(self):
+    def __init__(self, temperature=1.2):
         super(SoftCrossEntropyLoss, self).__init__()
+        self.temperature = temperature
 
     def forward(self, stu_outputs, tea_outputs):
-        input_log_likelihood = -F.log_softmax(stu_outputs, dim=1)
-        target_log_likelihood = F.softmax(tea_outputs, dim=1)
+        input_log_likelihood = -F.log_softmax(stu_outputs / self.temperature, dim=1)
+        target_log_likelihood = F.softmax(tea_outputs / self.temperature, dim=1)
         batch_size = stu_outputs.size(0)
         loss = torch.sum(torch.mul(input_log_likelihood, target_log_likelihood)) / batch_size
         return loss