From c3ca7d963e9ea02eed4d2890c5dc7f14e1882e8e Mon Sep 17 00:00:00 2001
From: 志浩 <neo.dzh@alibaba-inc.com>
Date: 星期四, 09 三月 2023 17:04:28 +0800
Subject: [PATCH] modify unit test for speech_diarization_sond-en-us-callhome-8k-n16k4-pytorch
---
funasr/tasks/diar.py | 10 ++++++++--
1 files changed, 8 insertions(+), 2 deletions(-)
diff --git a/funasr/tasks/diar.py b/funasr/tasks/diar.py
index 9a43945..45e3d7a 100644
--- a/funasr/tasks/diar.py
+++ b/funasr/tasks/diar.py
@@ -23,7 +23,8 @@
from funasr.layers.utterance_mvn import UtteranceMVN
from funasr.layers.label_aggregation import LabelAggregate
from funasr.models.ctc import CTC
-from funasr.models.encoder.resnet34_encoder import ResNet34Diar
+from funasr.models.encoder.resnet34_encoder import ResNet34Diar, ResNet34SpL2RegDiar
+from funasr.models.encoder.ecapa_tdnn_encoder import ECAPA_TDNN
from funasr.models.encoder.opennmt_encoders.conv_encoder import ConvEncoder
from funasr.models.encoder.opennmt_encoders.fsmn_encoder import FsmnEncoder
from funasr.models.encoder.opennmt_encoders.self_attention_encoder import SelfAttentionEncoder
@@ -121,10 +122,12 @@
fsmn=FsmnEncoder,
conv=ConvEncoder,
resnet34=ResNet34Diar,
+ resnet34_sp_l2reg=ResNet34SpL2RegDiar,
sanm_chunk_opt=SANMEncoderChunkOpt,
data2vec_encoder=Data2VecEncoder,
+ ecapa_tdnn=ECAPA_TDNN,
),
- type_check=AbsEncoder,
+ type_check=torch.nn.Module,
default="resnet34",
)
speaker_encoder_choices = ClassChoices(
@@ -158,6 +161,7 @@
classes=dict(
dot=DotScorer,
cosine=CosScorer,
+ conv=ConvEncoder,
),
type_check=torch.nn.Module,
default=None,
@@ -187,6 +191,8 @@
specaug_choices,
# --normalize and --normalize_conf
normalize_choices,
+ # --label_aggregator and --label_aggregator_conf
+ label_aggregator_choices,
# --model and --model_conf
model_choices,
# --encoder and --encoder_conf
--
Gitblit v1.9.1