Update modeling_wav2vec2_spkreg.py

modeling_wav2vec2_spkreg.py

@@ -529,6 +529,8 @@ class AMSoftmaxLoss(nn.Module):
         num_labels: int, 
         scale: float = 30.0, 
         margin: float = 0.35, 
+        label_smoothing: float = 0.0, 
+        reduction: str = "mean"
     ):
         """
         Args:
@@ -540,13 +542,13 @@ class AMSoftmaxLoss(nn.Module):
         self.num_labels = num_labels
         self.scale = scale
         self.margin = margin
+        self.label_smoothing = label_smoothing
+        self.reduction = reduction
 
     def forward(
         self, 
         inputs: torch.Tensor, 
         targets: torch.Tensor, 
-        label_smoothing: float = 0.0, 
-        reduction: str = "mean"
     ):
         """
         Args:
@@ -562,7 +564,9 @@ class AMSoftmaxLoss(nn.Module):
         psi = cosine - self.margin
         one_hot = nn.functional.one_hot(targets, self.num_labels)
         outputs = self.scale * torch.where(one_hot.bool(), psi, cosine)
-        loss = F.cross_entropy(outputs, targets, label_smoothing=label_smoothing, reduction=reduction)
+        loss = F.cross_entropy(
+            outputs, targets, label_smoothing=self.label_smoothing, reduction=self.reduction
+        )
         return loss
 
 
@@ -577,7 +581,9 @@ class AAMSoftmaxLoss(nn.Module):
         num_labels: int, 
         scale: float = 30.0, 
         margin: float = 0.35, 
-        easy_margin: bool = False
+        easy_margin: bool = False, 
+        label_smoothing: float = 0.0, 
+        reduction: str = "mean"
     ):
         """
         Args:
@@ -591,6 +597,8 @@ class AAMSoftmaxLoss(nn.Module):
         self.scale = scale
         self.margin = margin
         self.easy_margin = easy_margin
+        self.label_smoothing = label_smoothing
+        self.reduction = reduction
         
     def forward(
         self, 
@@ -627,7 +635,9 @@ class AAMSoftmaxLoss(nn.Module):
         outputs = (one_hot * phi) + ((1.0 - one_hot) * cosine)
         outputs = outputs * self.scale
 
-        loss = F.cross_entropy(outputs, targets, label_smoothing=label_smoothing, reduction=reduction)
+        loss = F.cross_entropy(
+            outputs, targets, label_smoothing=self.label_smoothing, reduction=self.reduction
+        )
         return loss
 
 
@@ -739,17 +749,24 @@ class Wav2Vec2SpkRegForSequenceClassification(Wav2Vec2SpkRegPreTrainedModel):
                 )
             elif self.config.loss_fct == 'additive_margin':
                 loss_fct = AMSoftmaxLoss(
-                    self.config.num_labels, self.config.scale, self.config.margin
+                    self.config.num_labels, 
+                    self.config.scale, 
+                    self.config.margin, 
+                    label_smoothing=self.config.label_smoothing, 
+                    reduction=self.config.reduction
                 )
             elif self.config.loss_fct == 'additive_angular_margin':
                 loss_fct = AAMSoftmaxLoss(
-                    self.config.num_labels, self.config.scale, self.config.margin, self.config.easy_margin
+                    self.config.num_labels, 
+                    self.config.scale, 
+                    self.config.margin, 
+                    self.config.easy_margin, 
+                    label_smoothing=self.config.label_smoothing, 
+                    reduction=self.config.reduction
                 )
             loss = loss_fct(
                 logits.view(-1, self.config.num_labels), 
                 labels.view(-1), 
-                label_smoothing=self.config.label_smoothing, 
-                reduction=self.config.reduction
             )
 
         if not return_dict: