Create speaker model config and it to registry

OML-Team · Dec 30, 2024 · 6d3fc3a · 6d3fc3a
1 parent a06b819
commit 6d3fc3a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/oml/configs/extractor/ecapa_tdnn_taoruijie.yaml b/oml/configs/extractor/ecapa_tdnn_taoruijie.yaml
@@ -0,0 +1,5 @@
+name: ecapa_tdnn
+args:
+  arch: ecapa_tdnn_taoruijie
+  weights: null
+  normalise_features: False
diff --git a/oml/registry/models.py b/oml/registry/models.py
@@ -4,6 +4,7 @@
 from torch import nn
 
 from oml.interfaces.models import IExtractor, IPairwiseModel
+from oml.models.audio.ecapa_tdnn.extractor import ECAPATDNNExtractor
 from oml.models.meta.projection import ExtractorWithMLP
 from oml.models.meta.siamese import (
     ConcatSiamese,
@@ -22,6 +23,7 @@
     "vit_clip": ViTCLIPExtractor,
     "vit_unicom": ViTUnicomExtractor,
     "extractor_with_mlp": ExtractorWithMLP,
+    "ecapa_tdnn": ECAPATDNNExtractor,
 }
 
 PAIRWISE_MODELS_REGISTRY = {