Upload model

config.json

@@ -1,4 +1,5 @@
 {
+  "_name_or_path": "facebook/wav2vec2-base",
   "activation_dropout": 0.0,
   "adapter_attn_dim": null,
   "adapter_kernel_size": 3,
@@ -6,11 +7,12 @@
   "add_adapter": false,
   "apply_spec_augment": true,
   "architectures": [
-    "Wav2Vec2ForPreTraining"
+    "Wav2Vec2SpkRegModel"
   ],
   "attention_dropout": 0.1,
   "auto_map": {
-    "AutoConfig": "configuration_wav2vec2_spkreg.Wav2Vec2SpkRegConfig"
+    "AutoConfig": "configuration_wav2vec2_spkreg.Wav2Vec2SpkRegConfig",
+    "AutoModel": "modeling_wav2vec2_spkreg.Wav2Vec2SpkRegModel"
   },
   "bos_token_id": 1,
   "classifier_proj_size": 256,
@@ -56,7 +58,6 @@
   "feat_quantizer_dropout": 0.0,
   "final_dropout": 0.0,
   "freeze_feat_extract_train": true,
-  "gradient_checkpointing": true,
   "hidden_act": "gelu",
   "hidden_dropout": 0.1,
   "hidden_size": 768,
@@ -119,6 +120,7 @@
     1,
     1
   ],
+  "torch_dtype": "float32",
   "transformers_version": "4.46.2",
   "use_weighted_layer_sum": false,
   "vocab_size": 32,

modeling_wav2vec2_spkreg.py

@@ -596,6 +596,9 @@ class AAMSoftmaxLoss(nn.Module):
         self.easy_margin = easy_margin
         self.label_smoothing = label_smoothing
         self.reduction = reduction
+
+        self.cos_m = math.cos(self.margin)
+        self.sin_m = math.sin(self.margin)
         
     def forward(
         self, 
@@ -611,9 +614,9 @@ 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)
+        cos_theta = torch.clamp(inputs, -1.0, 1.0)
         sin_theta = torch.sqrt(1.0 - torch.pow(cos_theta, 2))
-        psi = cos_theta * math.cos(self.margin) - sin_theta * math.sin(self.margin) # cos(theta + m)
+        psi = cos_theta * self.cos_m - sin_theta * self.sin_m # cos(theta + m)
         # theta = torch.acos(cos_theta)
         # psi = torch.cos(theta + self.margin)
         one_hot = nn.functional.one_hot(targets, num_labels)