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

@@ -612,8 +612,10 @@ 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)
+        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)
+        # theta = torch.acos(cos_theta)
+        # psi = torch.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(