Update modeling_wavlm_spkreg.py

modeling_wavlm_spkreg.py

@@ -456,7 +456,7 @@ class AAMSoftmaxLoss(nn.Module):
     def __init__(
         self, 
         scale: float = 30.0, 
-        margin: float = 0.35, 
+        margin: float = 0.2, 
         easy_margin: bool = False, 
         label_smoothing: float = 0.0, 
         reduction: str = "mean"
@@ -489,9 +489,23 @@ class AAMSoftmaxLoss(nn.Module):
         """
         _, num_labels = inputs.shape
         # `inputs` are the outputs from AngularLinear()
-        cos_theta = torch.clamp(inputs, -1.0 + 1e-7, 1.0 - 1e-7)
-        theta = torch.acos(cos_theta)
-        psi = torch.cos(theta + self.margin)
+        epsilon = 1e-6
+        # theta = torch.acos(cos_theta)
+        # psi = torch.cos(theta + self.margin)
+        cos_theta = torch.clamp(inputs, -1.0 + epsilon, 1.0 - epsilon)
+        sin_theta = torch.sqrt(1.0 - torch.pow(cos_theta, 2))
+        sin_theta = torch.clamp(sin_theta, 0.0 + epsilon, 1.0 - epsilon)
+
+        cos_m = math.cos(self.margin)
+        sin_m = math.sin(self.margin)
+        psi = cos_theta * cos_m - sin_theta * sin_m # cos(theta + m)
+
+        if self.easy_margin:
+            psi = torch.where(cos_theta > 0, psi, cos_theta)
+        else:
+            # Make the function cos(theta+m) monotonic decreasing while theta in [0°, 180°]
+            psi = torch.where((cos_theta - math.cos(math.pi - self.margin)) > 0, psi, cos_theta - self.margin)
+
         one_hot = nn.functional.one_hot(targets, num_labels)
         outputs = self.scale * torch.where(one_hot.bool(), psi, cos_theta)
         loss = F.cross_entropy(