app.py

SpeakerNet.py

@@ -0,0 +1,283 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy, sys, random
+from DatasetLoader import test_dataset_loader
+import importlib
+import time, itertools
+from utils.log import init_log
+from tqdm import tqdm
+import wandb
+from tuneThreshold import *
+
+
+class SpeakerNet(nn.Module):
+
+    def __init__(self, model, trainfunc, nPerSpeaker):
+        super(SpeakerNet, self).__init__()
+
+        self.model = model
+        self.loss = trainfunc
+        self.nPerSpeaker = nPerSpeaker
+
+    def forward(self, data, label=None):
+
+        data = data.reshape(-1, data.size()[-1])
+        outp = self.model(data)
+
+        if label == None:
+            return outp
+
+        else:
+
+            emb = outp.reshape(-1, self.nPerSpeaker, outp.size()[-1]).squeeze(1)
+            nloss, prec1 = self.loss(emb, label)
+
+
+            return nloss, prec1
+
+
+class Trainer(object):
+
+    def __init__(self, cfg, model, optimizer, scheduler, device):
+        self.cfg = cfg
+        self.model = model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.device = device
+        logging = init_log(cfg.save_dir)
+        self._print = logging.info
+        self.best = 0
+        self.test_eer = 0
+        self.test_mindcf = 0
+        self.best_model = []
+
+    def train(self, epoch, dataloader):
+        self.model.train()
+        pbar = tqdm(dataloader)
+        loss = 0
+        top1 = 0
+        index = 0
+        counter = 0
+
+        for data in pbar:
+            x, label = data[0].to(self.device), data[1].long().to(self.device)
+            nloss, prec1 = self.model(x, label)
+
+            self.optimizer.zero_grad()
+            nloss.backward()
+            self.optimizer.step()
+            # self.scheduler.step()
+
+            loss += nloss.detach().cpu().item()
+            top1 += prec1.detach().cpu().item()
+            index += x.size(0)
+            counter += 1
+
+            if self.cfg.wandb:
+                wandb.log({
+                    "epoch": epoch,
+                    "train_acc": top1 / counter,
+                    "train_loss": loss / counter,
+                })
+            pbar.set_description("Train Epoch:%3d ,Tloss:%.3f, Tacc:%.3f" % (epoch, loss/counter, top1/counter))
+
+        # self.scheduler.step()
+        self._print('epoch:{} - train loss: {:.3f} and train acc: {:.3f} total sample: {}'.format(
+            epoch, loss/counter, top1/counter, index))
+
+    def test(self, epoch, test_list, test_path, nDataLoaderThread, eval_frames, num_eval=10):
+
+        self.model.eval()
+        feats = {}
+
+        # read all lines
+        with open(test_list) as f:
+            lines = f.readlines()
+        files = list(itertools.chain(*[x.strip().split()[-2:] for x in lines]))
+        setfiles = list(set(files))
+        setfiles.sort()
+
+        # Define test data loader
+        test_dataset = test_dataset_loader(setfiles, test_path, eval_frames=eval_frames, num_eval=num_eval)
+
+        test_loader = torch.utils.data.DataLoader(
+            test_dataset,
+            batch_size=1,
+            shuffle=False,
+            num_workers=nDataLoaderThread,
+            drop_last=False,
+            sampler=None
+        )
+
+        # Extract features for every wav
+        for idx, data in enumerate(tqdm(test_loader)):
+
+            inp1 = data[0][0].to(self.device)  # (data[0]:(1,10,1024),data[1]:'id10270/GWXujl-xAVM/00017.wav')
+            with torch.no_grad():
+                ref_feat = self.model(inp1).detach().cpu()
+
+            feats[data[1][0]] = ref_feat
+
+        all_scores = []
+        all_labels = []
+        all_trials = []
+
+        # Read files and compute all scores
+        for idx, line in enumerate(tqdm(lines)):
+
+            data = line.split()
+
+            # Append random label if missing
+            if len(data) == 2:
+                data = [random.randint(0, 1)] + data
+            ref_feat = feats[data[1]].to(self.device)
+            com_feat = feats[data[2]].to(self.device)
+
+            if self.model.loss.test_normalize:
+                ref_feat = F.normalize(ref_feat, p=2, dim=1)
+                com_feat = F.normalize(com_feat, p=2, dim=1)
+
+            # dist = F.pairwise_distance(ref_feat.unsqueeze(-1),
+            #                            com_feat.unsqueeze(-1).transpose(0, 2)).detach().cpu().numpy()
+            #
+            # score = -1 * numpy.mean(dist)
+            dist = F.cosine_similarity(ref_feat.unsqueeze(-1),
+                                       com_feat.unsqueeze(-1).transpose(0, 2)).detach().cpu().numpy()
+            score = numpy.mean(dist)
+
+            all_scores.append(score)
+            all_labels.append(int(data[0]))
+            all_trials.append(data[1] + " " + data[2])
+
+        result = tuneThresholdfromScore(all_scores, all_labels, [1, 0.1])
+        fnrs, fprs, thresholds = ComputeErrorRates(all_scores, all_labels)
+        mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, self.cfg.dcf_p_target, self.cfg.dcf_c_miss, self.cfg.dcf_c_fa)
+        self.test_eer = result[1]
+        self.test_mindcf = mindcf
+        self.threshold = threshold
+        if self.cfg.wandb:
+            wandb.log({
+                "test_eer": self.test_eer,
+                "test_MinDCF": self.test_mindcf,
+            })
+        self._print('epoch:{} - test EER: {:.3f} and test MinDCF: {:.3f} total sample: {} threshold: {:.3f}'.format(
+            epoch, self.test_eer, self.test_mindcf, len(lines), self.threshold))
+
+        return self.test_eer
+
+    def save_model(self, epoch):
+        if self.test_eer < self.best or self.best == 0:
+            self.best = self.test_eer
+            if self.cfg.wandb:
+                wandb.run.summary["best_accuracy"] = self.best
+            model_state_dict = self.model.state_dict()
+            optimizer_state_dict = self.optimizer.state_dict()
+            scheduler_state_dict = self.scheduler.state_dict()
+            file_save_path = 'epoch:%d,EER:%.4f,MinDCF:%.4f' % (epoch, self.test_eer, self.test_mindcf)
+            if not os.path.exists(self.cfg.save_dir):
+                os.mkdir(self.cfg.save_dir)
+            torch.save({
+                'epoch': epoch,
+                'test_eer':  self.test_eer,
+                'test_mindcf': self.test_mindcf,
+                'model_state_dict': model_state_dict,
+                'optimizer_state_dict': optimizer_state_dict,
+                'scheduler_state_dict': scheduler_state_dict},
+                os.path.join(self.cfg.save_dir, file_save_path))
+            self.best_model.append(file_save_path)
+            if len(self.best_model) > 3:
+                del_file = os.path.join(self.cfg.save_dir, self.best_model.pop(0))
+                if os.path.exists(del_file):
+                    os.remove(del_file)
+                else:
+                    print("no exists {}".format(del_file))
+        # 每20个epoch保存一下
+        if epoch % 20 == 0:
+            model_state_dict = self.model.state_dict()
+            optimizer_state_dict = self.optimizer.state_dict()
+            scheduler_state_dict = self.scheduler.state_dict()
+            file_save_path = 'epoch:%d,EER:%.4f,MinDCF:%.4f' % (epoch, self.test_eer, self.test_mindcf)
+            if not os.path.exists(self.cfg.save_dir):
+                os.mkdir(self.cfg.save_dir)
+            if not os.path.exists(os.path.join(self.cfg.save_dir, file_save_path)):
+                torch.save({
+                    'epoch': epoch,
+                    'test_eee':  self.test_eer,
+                    'test_mindcf': self.test_mindcf,
+                    'model_state_dict': model_state_dict,
+                    'optimizer_state_dict': optimizer_state_dict,
+                    'scheduler_state_dict': scheduler_state_dict},
+                    os.path.join(self.cfg.save_dir, file_save_path))
+
+    def scoretxt(self, score_file, test_list, test_path, eval_frames, num_eval=10):
+
+        self.model.eval()
+        feats = {}
+
+        # read all lines
+        with open(test_list) as f:
+            lines = f.readlines()
+        files = list(itertools.chain(*[x.strip().split()[-2:] for x in lines]))
+        setfiles = list(set(files))
+        setfiles.sort()
+
+        # Define test data loader
+        test_dataset = test_dataset_loader(setfiles, test_path, eval_frames=eval_frames, num_eval=num_eval)
+
+        test_loader = torch.utils.data.DataLoader(
+            test_dataset,
+            batch_size=1,
+            shuffle=False,
+            drop_last=False,
+            sampler=None
+        )
+
+        # Extract features for every wav
+        for idx, data in enumerate(tqdm(test_loader)):
+
+            inp1 = data[0][0].to(self.device)  # (data[0]:(1,10,1024),data[1]:'id10270/GWXujl-xAVM/00017.wav')
+            with torch.no_grad():
+                ref_feat = self.model(inp1).detach().cpu()
+
+            feats[data[1][0]] = ref_feat
+
+
+        f = open(score_file, "w")
+        # Read files and compute all scores
+        for idx, line in enumerate(tqdm(lines)):
+
+            data = line.split()
+
+            # Append random label if missing
+            ref_feat = feats[data[-2]].to(self.device)
+            com_feat = feats[data[-1]].to(self.device)
+
+            if self.model.loss.test_normalize:
+                ref_feat = F.normalize(ref_feat, p=2, dim=1)
+                com_feat = F.normalize(com_feat, p=2, dim=1)
+
+            # dist = F.pairwise_distance(ref_feat.unsqueeze(-1),
+            #                            com_feat.unsqueeze(-1).transpose(0, 2)).detach().cpu().numpy()
+            #
+            # score = -1 * numpy.mean(dist)
+            dist = F.cosine_similarity(ref_feat.unsqueeze(-1),
+                                       com_feat.unsqueeze(-1).transpose(0, 2)).detach().cpu().numpy()
+            score = numpy.mean(dist)
+
+            score_line = str(score) + " " + data[-2] + " " + data[-1]
+            f.write(score_line+'\n')
+        f.close()
+
+
+
+
+
+
+
+
+
+
+
+

app.py

@@ -0,0 +1,77 @@
+import gradio as gr
+import torch
+import net
+import argparse
+from config import set_cfg, cfg
+from SpeakerNet import *
+import lossfunction
+from DatasetLoader import loadWAV
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--config_name", type=str, default="ECAPA_TDNN_data_cfg", help="the configs name that will as a base configs")
+parser.add_argument("--resume", default="train_models/epoch_37_ECAPA_TDNN2.48", type=str, help="resume path")
+args = parser.parse_args()
+global cfg
+assert args.config_name is not None
+if args.config_name:
+    set_cfg(args.config_name)
+cfg.replace(vars(args))
+del args
+
+device = torch.device("cpu")
+model = getattr(net, cfg.model)().to(device)
+loss = getattr(lossfunction, cfg.loss)(cfg.nOut, cfg.nClasses).to(device)
+model = SpeakerNet(model=model, trainfunc=loss, nPerSpeaker=cfg.nPerSpeaker)
+
+ckpt = torch.load("train_models/epoch_37_ECAPA_TDNN2.48", map_location="cpu")
+model.load_state_dict(ckpt['model_state_dict'], strict=False)
+print("checkpoint加载完毕!")
+
+model.eval()
+
+def SpeakerVerification(path1,path2):
+    inp1 = loadWAV(path1, max_frames=300, evalmode=True)
+    inp2 = loadWAV(path2, max_frames=300, evalmode=True)
+    # print(inp1.shape)
+    inp1 = torch.FloatTensor(inp1)
+    inp2 = torch.FloatTensor(inp2)
+    # print(inp1.shape)
+    with torch.no_grad():
+        emb1 = model(inp1).detach().cpu()
+        emb2 = model(inp2).detach().cpu()
+    emb1 = F.normalize(emb1, p=2, dim=1)
+    emb2 = F.normalize(emb2, p=2, dim=1)
+    dist = F.cosine_similarity(emb1.unsqueeze(-1),  emb2.unsqueeze(-1).transpose(0, 2)).numpy()
+    score = numpy.mean(dist)
+    print(score)
+    # threshold = 0.414
+    if score >= 0.414:
+        output = "同一个人"
+    else:
+        output = "不同的人"
+
+    return output
+
+inputs = [
+    gr.inputs.Audio(source="upload", type="filepath", label="Speaker #1", optional=False),
+    gr.inputs.Audio(source="upload", type="filepath", label="Speaker #2", optional=False)
+]
+
+
+examples = [["example/speaker1-1.wav", "example/speaker1-2.wav"],
+            ["example/speaker1-1.wav", "example/speaker2-1.wav"],
+            ["example/speaker2-1.wav", "example/speaker2-2.wav"],
+            ["example/speaker1-2.wav", "example/speaker2-2.wav"],
+            ["example/speaker3-1.wav", "example/speaker3-2.wav"],
+            ["example/speaker3-1.wav", "example/speaker4-1.wav"],
+            ["example/speaker4-1.wav", "example/speaker4-2.wav"],
+            ["example/speaker3-2.wav", "example/speaker4-2.wav"],
+            ["example/speaker4-1.wav", "example/speaker5-2.wav"],
+            ]
+
+iface = gr.Interface(fn=SpeakerVerification, inputs=inputs, outputs="text", examples=examples)
+iface.launch(share=True)
+
+if __name__ == '__main__':
+    # print(SpeakerVerification("example/speaker1-1.wav", "example/speaker1-2.wav"))
+    pass

config.py

@@ -0,0 +1,517 @@
+class Config(object):
+    def __init__(self, config_dict: dict):
+        for key, val in config_dict.items():
+            if val is not None:
+                self.__setattr__(key, val)
+
+    def copy(self, new_config_dict={}):
+        ret = Config(vars(self))
+        for key, val in new_config_dict.items():
+            if val is not None:
+                ret.__setattr__(key, val)
+        return ret
+
+    def replace(self, new_config_dict):
+        if isinstance(new_config_dict, Config):
+            new_config_dict = vars(new_config_dict)
+
+        for key, val in new_config_dict.items():
+            if val is not None:
+                self.__setattr__(key, val)
+
+    def print(self):
+        for k, v in vars(self).items():
+            print(k, '=', v)
+
+    # def parser_val(self, val):
+    #     if isinstance(val, dict):
+    #         return Config(val)
+    #     elif isinstance(val, list):
+    #         for i in range(len(val)):
+    #             if val is not None:
+    #                 val[i] = self.parser_val(val[i])
+    #         return val
+    #     else:
+    #         return val
+
+    def __str__(self):
+        return str(vars(self))
+
+
+base_config = Config({
+    "project": "speaker_verification",
+    "name": "VGGVox",
+    "save_dir": "train_models/",
+    "resume": "",
+
+    # Training and test data
+    "dataset": Config({
+            "name": "voxceleb2_wav",
+            "train_list": "data/train_list.txt",
+            "test_list": "data/veri_list.txt",
+            "train_path": "data/voxceleb2",
+            "test_path": "data/voxceleb1",
+            "musan_path": "data/musan_split",  # 噪声文件
+            "rir_path": "data/RIRS_NOISES/simulated_rirs",  # 混响文件
+        }),
+
+
+    # Data loader
+    "max_frames": 300,  # 训练时帧长
+    "eval_frames": 300,
+    "batch_size": 64,
+    "max_seg_per_spk": 500,  # 每个说话人最大的语音段数
+    "nDataLoaderThread": 16,  # 多线程加载
+    "augment": True,  # 是否数据增强
+    "seed": 10,
+    "segment": 1,
+
+    # Training details
+    "test_interval": 1,  # 测试间隔
+    "max_epoch": 500,
+
+    # Model definition
+    "n_mels": 40,
+    "log_input": False,
+    "model": "Vgg",
+    "encoder_type": "SAP",
+    "nOut": 512,
+
+    # Loss functions
+    "loss": "SoftmaxProto",  # lossfunction function
+    "hard_prob": 0.5,
+    "hard_rank": 10,
+    "margin": 0.2,
+    "scale": 30,
+    "nPerSpeaker": 2,  # 同一段语音取多少组
+    "nClasses": 5994,
+
+    # Optimizer
+    "optimizer": "adam",
+    "scheduler": "steplr",
+    "lr": 0.001,
+    "lr_decay": 0.95,
+    "weight_decay": 0,
+
+    # Evaluation parameters
+    "dcf_p_target": 0.05,
+    "dcf_c_miss": 1,
+    "dcf_c_fa": 1,
+
+    # eval
+    "eval": False,
+})
+
+cfg = base_config
+
+vgg_cfg = Config({
+    "name": "vgg_spectrogram",
+    "model": "vgg",
+    "batch_size": 64,
+    "nPerSpeaker": 2,
+})
+
+Unet_cfg = Config({
+    "name": "Unet",
+    "model": "UNetVgg",
+    "batch_size": 48,
+    "nPerSpeaker": 2,
+    "loss": "Unetloss"
+})
+
+UnetMask_cfg = Config({
+    "name": "UnetMask",
+    "model": "UNetVggMask",
+    "batch_size": 16,
+    "nPerSpeaker": 2,
+    "segment": 3,
+    "loss": "UnetMaskloss"
+})
+
+ECAPA_TDNN_cfg = Config({
+    "name": "ECAPA_TDNNm",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmaxProto",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+})
+
+ECAPA_TDNNm_cfg = Config({
+    "name": "ECAPA_TDNNm",
+    "model": "ECAPA_TDNN",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+})
+
+ECAPA_TDNN1024_cfg = Config({
+    "name": "ECAPA_TDNN1024",
+    "model": "ECAPA_TDNN",
+    "batch_size": 80,
+    "nPerSpeaker": 2,
+    "channels": 1024,
+    "nOut": 192,
+})
+
+ECAPA_TDNN_ks5_cfg = Config({
+    "name": "ECAPA_TDNN_ks5",
+    "model": "ECAPA_TDNN_ks5",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+})
+
+ECAPA_TDNN_L2_cfg = Config({
+    "name": "ECAPA_TDNN_L2_pre",
+    "model": "ECAPA_TDNN_L2",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+    "resume": "train_models/speaker_verification_ECAPA_TDNN/20210725/epoch:47,EER:2.5981,MinDCF:0.1912"
+})
+
+ECAPA_TDNN_br_cfg = Config({
+    "name": "ECAPA_TDNN_br",
+    "model": "ECAPA_TDNN_br",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+})
+
+ECAPATDNN_cfg = Config({
+    "name": "ECAPATDNN",
+    "model": "ECAPATDNN",
+    "batch_size": 110,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+    "input_size": 80,
+})
+
+HRNet_cfg = Config({
+    "name": "hrnet",
+    "model": "hrnet",
+    "max_frames": 224,
+    "eval_frames": 224,
+    "batch_size": 48,
+    "nPerSpeaker": 2,
+    "nOut": 1024,
+    "input_size": 224*224,
+
+    "model_cfg": Config({
+        "hrnet_name": "w48",
+        "STAGE1": {
+            "NUM_MODULES": 1,
+            "NUM_BRANCHES": 1,
+            "BLOCK": "BOTTLENECK",
+            "NUM_BLOCKS": [4],
+            "NUM_CHANNELS": [64],
+            "FUSE_METHOD": "SUM"
+        },
+        "STAGE2": {
+            "NUM_MODULES": 1,
+            "NUM_BRANCHES": 2,
+            "BLOCK": "BASIC",
+            "NUM_BLOCKS": [4, 4],
+            "NUM_CHANNELS": [18, 36],
+            "FUSE_METHOD": "SUM"
+        },
+        "STAGE3": {
+            "NUM_MODULES": 4,
+            "NUM_BRANCHES": 3,
+            "BLOCK": "BASIC",
+            "NUM_BLOCKS": [4, 4, 4],
+            "NUM_CHANNELS": [18, 36, 72],
+            "FUSE_METHOD": "SUM"
+        },
+        "STAGE4": {
+            "NUM_MODULES": 3,
+            "NUM_BRANCHES": 4,
+            "BLOCK": "BASIC",
+            "NUM_BLOCKS": [4, 4, 4, 4],
+            "NUM_CHANNELS": [18, 36, 72, 144],
+            "FUSE_METHOD": "SUM"
+        },
+    }),
+
+})
+
+VGG_TDNN_cfg = Config({
+    "name": "Vggtdnn1",
+    "model": "Vggtdnn",
+    "batch_size": 256,
+    "nOut": 512,
+    "nDataLoaderThread": 16,
+})
+
+ResNetSE34V2_cfg = Config({
+    "name": "ResNetSE34V2",
+    "model": "ResNetSE34V2",
+    "batch_size": 128,
+    "nOut": 512,
+    "nDataLoaderThread": 16,
+})
+
+HRTDNN_cfg = Config({
+    "name": "hrtdnn",
+    "model": "hrtdnn",
+    "max_frames": 300,
+    "eval_frames": 300,
+    "batch_size": 96,
+    "nPerSpeaker": 2,
+    "nOut": 256,
+
+    "model_cfg": Config({
+        "hrnet_name": "hrtdnn",
+        "STAGE1": {
+            "NUM_BRANCHES": 1,
+            "BLOCK": 'TDNNBlock',
+            "NUM_CHANNELS": [128],
+            "FUSE_METHOD": "SUM"
+        },
+        "STAGE2": {
+            "NUM_BRANCHES": 2,
+            "BLOCK": 'TDNNBlock',
+            "NUM_CHANNELS": [128, 512],
+            "FUSE_METHOD": "SUM"
+        },
+        "STAGE3": {
+            "NUM_BRANCHES": 3,
+            "BLOCK": 'TDNNBlock',
+            "NUM_CHANNELS": [128, 512, 1024],
+            "FUSE_METHOD": "SUM"
+        },
+
+    }),
+
+})
+
+ResTDNN_cfg = Config({
+    "name": "ResTDNN",
+    "model": "ResTDNN",
+    "batch_size": 110,
+    "nOut": 256,
+    "nDataLoaderThread": 16,
+})
+
+TDNN_VGG_cfg = Config({
+    "name": "TDNN_VGG",
+    "model": "TDNN_VGG",
+    "batch_size": 64,
+    "nOut": 256,
+    "nDataLoaderThread": 16,
+})
+
+ResNet_TDNN_cfg = Config({
+    "name": "ResNet_TDNN",
+    "model": "ResNet_TDNN",
+    "batch_size": 96,
+    "nOut": 192,
+    "nDataLoaderThread": 16,
+})
+
+ResNet_TDNNa_cfg = Config({
+    "name": "ResNet_TDNNa",
+    "model": "ResNet_TDNN",
+    "batch_size": 96,
+    "nOut": 192,
+    "nDataLoaderThread": 16,
+})
+
+ResNet_TDNNaam_cfg = Config({
+    "name": "ResNet_TDNNaam",
+    "model": "ResNet_TDNN",
+    "loss": "AamSoftmaxProto",
+    "margin": 0.2,
+    "scale": 30,
+    "batch_size": 96,
+    "nOut": 192,
+    "nDataLoaderThread": 16,
+    "augment": True,
+})
+
+TDNN_ResNet_cfg = Config({
+    "name": "TDNN_ResNet",
+    "model": "TDNN_ResNet",
+    "batch_size": 48,
+    "nOut": 256,
+    "nDataLoaderThread": 16,
+})
+
+hr_tdnn_cfg = Config({
+    "name": "hr_tdnn",
+    "model": "hr_tdnn",
+    "batch_size": 46,
+    "nOut": 192,
+    "nDataLoaderThread": 16,
+})
+
+
+ECAPA_TDNNma_cfg = Config({
+    "name": "ECAPA_TDNNma",
+    "model": "ECAPA_TDNN",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+    "augment": True,
+})
+
+ECAPA_TDNNaam_cfg = Config({
+    "name": "ECAPA_TDNNaam",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+})
+
+ECAPA_TDNNaam1_cfg = Config({
+    "name": "ECAPA_TDNNaam1",
+    "model": "ECAPA_TDNN",
+    "loss": "AdditiveAngularMargin",
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+})
+
+ECAPA_TDNNaam2_cfg = Config({
+    "name": "ECAPA_TDNNaam2",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "margin": 0.2,
+    "scale": 30,
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+
+})
+
+ECAPA_TDNNaam3_cfg = Config({
+    "name": "ECAPA_TDNNaam3",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "margin": 0.1,
+    "scale": 30,
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+
+})
+
+ECAPA_TDNN_aamproto_cfg = Config({
+    "name": "ECAPA_TDNN_aamproto",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmaxProto",
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+    "augment": True,
+})
+
+ECAPA_TDNN_aamproto1_cfg = Config({
+    "name": "ECAPA_TDNN_aamproto1",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmaxProto",
+    "margin": 0.2,
+    "scale": 30,
+    "batch_size": 180,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+    "augment": True,
+})
+
+ECAPA_TDNN0_cfg = Config({
+    "name": "ECAPA_TDNN-1lr0.001",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "batch_size": 360,
+    "nOut": 192,
+    "nPerSpeaker": 1,
+    "resume": "train_models/speaker_verification_ECAPA_TDNN0/20210928/epoch:25,EER:2.4125,MinDCF:0.1537",
+})
+
+SwinTransformer_cfg = Config({
+    "name": "SwinTransformer",
+    "model": "SwinTransformer",
+    "loss": "SoftmaxProto",
+    "max_frames": 224,
+    "eval_frames": 224,
+    "n_mels": 224,
+    "batch_size": 90,
+    "nPerSpeaker": 2,
+    "nOut": 192,
+    "augment": True,
+    "lr": 5e-5,
+})
+
+ECAPA_TDNN_aampre_cfg = Config({
+    "name": "ECAPA_TDNN_aampre",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmaxProto",
+    "batch_size": 180,
+    "nOut": 192,
+    "nPerSpeaker": 2,
+    "resume": "train_models/speaker_verification_ECAPA_TDNNma/20210908/epoch:67,EER:2.3224,MinDCF:0.1658",
+})
+
+# 更换dataloader
+ECAPA_TDNN_data_cfg = Config({
+    "name": "ECAPA_TDNN_data",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+
+})
+
+# 标准的ECAPA_TDNN 学习率CyclicLR
+ECAPA_TDNNaam_cyclr_cfg = Config({
+    "name": "ECAPA_TDNNaam_cyclr",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "margin": 0.2,
+    "scale": 30,
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+
+})
+
+# 跟换数据加载的ResNet_TDNN只用softmax
+ResNet_TDNNaam_data_cfg = Config({
+    "name": "ResNet_TDNNaam_data",
+    "model": "ResNet_TDNN",
+    "loss": "AamSoftmax",
+    "margin": 0.2,
+    "scale": 30,
+    "batch_size": 192,
+    "nOut": 192,
+    "nDataLoaderThread": 16,
+    "nPerSpeaker": 1,
+    "augment": True,
+})
+
+# 更换dataloader, 和cyclical lr
+ECAPA_TDNN_dataClr_cfg = Config({
+    "name": "ECAPA_TDNN_dataClr",
+    "model": "ECAPA_TDNN",
+    "loss": "AamSoftmax",
+    "batch_size": 360,
+    "nPerSpeaker": 1,
+    "nOut": 192,
+    "augment": True,
+})
+
+def set_cfg(config_name: str):
+    """ Sets the active configs. Works even if cfg is already imported! """
+    global cfg
+    # Note this is not just an eval because I'm lazy, but also because it can
+    # be used like ssd300_config.copy({'max_size': 400}) for extreme fine-tuning
+    cfg.replace(eval(config_name))

dataloader.py

@@ -0,0 +1,75 @@
+import torch
+from torch.utils.data import Dataset
+from DatasetLoader import AugmentWAV, loadWAV
+import os
+import numpy as np
+import random
+
+
+class TrainDataset(Dataset):
+    def __init__(self, train_list, train_path, augment, musan_path, rir_path, max_frames,):
+        self.train_list = train_list
+        self.max_frames = max_frames
+        self.augment_wav = AugmentWAV(musan_path=musan_path, rir_path=rir_path, max_frames=max_frames)
+        self.augment = augment
+        self.musan_path = musan_path
+        self.rir_path = rir_path
+
+        with open(train_list) as dataset_file:
+            lines = dataset_file.readlines()
+
+        dictkeys = list(set([x.split()[0] for x in lines]))
+        dictkeys.sort()
+        dictkeys = {key: ii for ii, key in enumerate(dictkeys)}
+
+        np.random.seed(100)
+        np.random.shuffle(lines)
+
+        self.data_list = []
+        self.data_label = []
+
+        for lidx, line in enumerate(lines):
+            data = line.strip().split()
+            speaker_label = dictkeys[data[0]]
+            filename = os.path.join(train_path, data[1])
+
+            self.data_list.append(filename)
+            self.data_label.append(speaker_label)
+
+    def __getitem__(self, index):
+
+        audio = loadWAV(self.data_list[index], self.max_frames, evalmode=False)
+        if self.augment:
+            augtype = random.randint(0, 4)  # 包括0，4
+            if augtype == 1:
+                audio = self.augment_wav.reverberate(audio)
+            elif augtype == 2:
+                audio = self.augment_wav.additive_noise('music', audio)
+            elif augtype == 3:
+                audio = self.augment_wav.additive_noise('speech', audio)
+            elif augtype == 4:
+                audio = self.augment_wav.additive_noise('noise', audio)
+
+        return torch.FloatTensor(audio), self.data_label[index]
+
+    def __len__(self):
+        return len(self.data_list)
+
+
+if __name__ == "__main__":
+    train_dataset = TrainDataset(train_list="data/train_list.txt", augment=True,
+                                  musan_path="data/musan_split", rir_path="data/RIRS_NOISES/simulated_rirs",
+                                  max_frames=300, train_path="data/voxceleb2")
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=32,
+        pin_memory=False,
+        drop_last=True,
+    )
+    x, y = iter(train_loader).next()
+    print("x:", x.shape, "y:", y.shape)
+
+
+
+
+

requirements.txt

@@ -0,0 +1 @@
+wandb

train.py

@@ -0,0 +1,186 @@
+from torch import optim
+import argparse
+from datetime import datetime
+import wandb
+import torch.backends.cudnn as cudnn
+from torch import optim
+from torch.utils.data import DataLoader
+from torchinfo import summary
+from timm.scheduler.cosine_lr import CosineLRScheduler
+
+import lossfunction
+import net
+from DatasetLoader import *
+from dataloader import TrainDataset
+from SpeakerNet import *
+from config import set_cfg, cfg
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--config_name", type=str, default="", help="the configs name that will as a base configs")
+    parser.add_argument("--project", default=None, type=str, help="project name")
+    parser.add_argument("--name", default=None, type=str, help="run name")
+    parser.add_argument("--save_dir", default=None, type=str, help="save path")
+    parser.add_argument("--resume", default=None, type=str, help="resume path")
+    parser.add_argument("--dataset", default=None, type=str, help="dataset path")
+
+    parser.add_argument("--epoch", default=None, type=int, help="max epoch")
+    parser.add_argument("--test_freq", default=None, type=int, help="frequency test epoch")
+    parser.add_argument("--batch_size", default=None, type=int, help="batch size")
+    parser.add_argument("--lr", default=None, type=float, help="learning rate")
+    parser.add_argument("--seed", default=None, type=int)
+
+    parser.add_argument("--wandb", action='store_true', default=False, help='use wandb to log ')
+    parser.add_argument("--note", type=str, default="", help='wandb note')
+
+    parser.add_argument('--eval', dest='eval', action='store_true', default=False, help='Eval only')
+    parser.add_argument('--score', dest='score', action='store_true', default=False, help='Eval only')
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    global cfg
+    args = get_args()
+    assert args.config_name is not None
+    if args.config_name:
+        set_cfg(args.config_name)
+    cfg.replace(vars(args))
+    del args
+
+    cfg.save_dir = os.path.join(cfg.save_dir, cfg.project + '_' + cfg.name, datetime.now().strftime('%Y%m%d'))
+    if not os.path.exists(cfg.save_dir):
+        os.makedirs(cfg.save_dir)
+
+    if cfg.wandb:
+        wandb.login(host="http://49.233.11.7:8080", key="local-7dc64cc63778f0723dc202d2624a97cef7043120")
+        wandb.init(project=cfg.project, name=cfg.name, config=cfg, save_code=True, notes=cfg.note)
+
+    # cudnn related setting
+    cudnn.benchmark = True
+    torch.backends.cudnn.deterministic = False
+    torch.backends.cudnn.enabled = True
+
+    start_epoch = 1
+
+    # ---------------模型---------------
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # device = torch.device("cpu")
+    # model = getattr(net, cfg.model)(cfg.nOut, cfg.encoder_type, cfg.log_input).to(device)
+    #  ------ECAPA_TDNN.yaml------ResNet_TDNN----
+    model = getattr(net, cfg.model)().to(device)
+
+    # loss = getattr(lossfunction, cfg.loss)(cfg.nOut, cfg.nClasses, cfg.margin, cfg.scale).to(device)
+    # ----aamsoftmax----
+    loss = getattr(lossfunction, cfg.loss)(cfg.nOut, cfg.nClasses).to(device)
+
+    # model = SpeakerUnet(model=model, trainfunc=loss, nPerSpeaker=cfg.nPerSpeaker, segment=cfg.segment)
+    model = SpeakerNet(model=model, trainfunc=loss, nPerSpeaker=cfg.nPerSpeaker)
+    # swin
+    optimizer = optim.AdamW(model.parameters(), eps=1e-8, betas=(0.9, 0.999),
+                lr=cfg.lr, weight_decay=0.05)
+    # optimizer = optim.Adam(model.parameters(), lr=cfg.lr, weight_decay=0.000002)
+    # scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 50, 70], gamma=0.1, last_epoch=-1)
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5,
+                                                     threshold=0.001, threshold_mode='rel',
+                                                     cooldown=0, min_lr=1e-5, eps=1e-08, verbose=True)
+    # scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=cfg.lr, max_lr=0.003, mode='triangular2',
+    #                                         step_size_up=12000, cycle_momentum=False)
+
+
+    if cfg.resume:
+        # ckpt = torch.load(cfg.resume, map_location="cpu")
+        ckpt = torch.load(cfg.resume)
+        model.load_state_dict(ckpt['model_state_dict'], strict=False)
+        # optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+        # scheduler.load_state_dict(ckpt['scheduler_state_dict'])
+        # start_epoch = ckpt['epoch'] + 1
+        print("checkpoint加载完毕!")
+    # print(model)
+
+    # test, eval, train
+    trainer = Trainer(cfg, model, optimizer, scheduler, device)
+
+    # ---------------score--------------
+    if cfg.score:
+        score_dir = os.path.join('score', cfg.name+"_"+datetime.now().strftime('%Y%m%d'))
+        if not os.path.exists(score_dir):
+            os.makedirs(score_dir)
+        score_file = os.path.join(score_dir, 'scores.txt')
+        trainer.scoretxt(score_file, 'data/voxsrc2021_blind.txt', 'data/voxsrc2021', cfg.eval_frames)
+        # trainer.scoretxt(score_file, cfg.dataset.test_list, cfg.dataset.test_path, cfg.eval_frames)
+    # ---------------eval--------------
+    elif cfg.eval:
+        trainer.test(0, cfg.dataset.test_list, cfg.dataset.test_path, cfg.nDataLoaderThread, cfg.eval_frames)
+    else:
+        # ---------------训练--------------
+        train_dataset = train_dataset_loader(train_list=cfg.dataset.train_list,
+                                             augment=cfg.augment, musan_path=cfg.dataset.musan_path,
+                                             rir_path=cfg.dataset.rir_path, max_frames=cfg.max_frames,
+                                             segment=cfg.segment, train_path=cfg.dataset.train_path)
+
+        train_sampler = train_dataset_sampler(train_dataset, nPerSpeaker=cfg.nPerSpeaker,
+                                              max_seg_per_spk=cfg.max_seg_per_spk, batch_size=cfg.batch_size,
+                                              seed=cfg.seed)
+
+        # train_dataset = TrainDataset(train_list=cfg.dataset.train_list,
+        #                              augment=cfg.augment, musan_path=cfg.dataset.musan_path,
+        #                              rir_path=cfg.dataset.rir_path, max_frames=cfg.max_frames,
+        #                              train_path=cfg.dataset.train_path)
+
+        train_loader = torch.utils.data.DataLoader(
+            train_dataset,
+            batch_size=cfg.batch_size,
+            num_workers=cfg.nDataLoaderThread,
+            sampler=train_sampler,
+            pin_memory=False,
+            drop_last=True,
+        )
+
+        x, y = iter(train_loader).next()
+        print('x.shape:', x.shape, 'y.shape:', y.shape)
+        print('x.dtype:', x.dtype, 'y.dtype:', y.dtype)
+
+        summary(model, input_size=(tuple(x.shape)))
+
+        it = 0
+        min_eer = float("inf")
+        for epoch in range(start_epoch, cfg.max_epoch):
+            trainer.train(epoch, train_loader)
+            if epoch % cfg.test_interval == 0:
+                eer = trainer.test(epoch, cfg.dataset.test_list, cfg.dataset.test_path, cfg.nDataLoaderThread,
+                                   cfg.eval_frames)
+                scheduler.step(eer)
+                # # -----Clr------
+                # if eer < min_eer:
+                #     min_eer = eer
+                #     it = 0
+                #
+                # else:
+                #     it += 1
+                #
+                #     if it >= 8:
+                #         lr = cfg.lr * 0.1
+                #         trainer.scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=lr, max_lr=cfg.lr,
+                #                                                 mode='triangular2',
+                #                                                 step_size_up=6000, cycle_momentum=False)
+                #         it = 0
+                # # -----Clr------
+                trainer.save_model(epoch)
+
+    print("finishing")
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+
+
+
+
+

tuneThreshold.py

@@ -0,0 +1,62 @@
+from sklearn import metrics
+import numpy
+from operator import itemgetter
+
+
+def tuneThresholdfromScore(scores, labels, target_fa, target_fr=None):
+    fpr, tpr, thresholds = metrics.roc_curve(labels, scores, pos_label=1)
+    fnr = 1 - tpr
+
+    tunedThreshold = []
+    if target_fr:
+        for tfr in target_fr:
+            idx = numpy.nanargmin(numpy.absolute((tfr - fnr)))
+            tunedThreshold.append([thresholds[idx], fpr[idx], fnr[idx]])
+
+    for tfa in target_fa:
+        idx = numpy.nanargmin(numpy.absolute((tfa - fpr)))  # numpy.where(fpr<=tfa)[0][-1] nanargmin 返回轴上最小的值忽略Nans
+        tunedThreshold.append([thresholds[idx], fpr[idx], fnr[idx]])
+
+    idxE = numpy.nanargmin(numpy.absolute((fnr - fpr)))
+    eer = max(fpr[idxE], fnr[idxE]) * 100
+
+    return tunedThreshold, eer, fpr, fnr
+
+# Creates a list of false-negative rates, a list of false-positive rates
+# and a list of decision thresholds that give those error-rates.
+def ComputeErrorRates(scores, labels):
+    sorted_indexes, thresholds = zip(*sorted([(index, threshold) for index, threshold in enumerate(scores)],
+                                           key=itemgetter(1)))
+    labels = [labels[i] for i in sorted_indexes]
+    fnrs = []  # 负样本接受
+    fprs = []  # 正样本接受
+
+    for i in range(0, len(labels)):
+        if i == 0:
+            fnrs.append(labels[i])
+            fprs.append(1 - labels[i])
+        else:
+            fnrs.append(fnrs[i-1] + labels[i])
+            fprs.append(fprs[i-1] + 1 - labels[i])
+
+    fnrs_norm = sum(labels)  # 真正样本个数
+    fprs_norm = len(labels) - fnrs_norm  # 负样本个数
+    fnrs = [x / float(fnrs_norm) for x in fnrs]  # 错误的拒绝 正样本分错的比例
+    fprs = [1 - x / float(fprs_norm) for x in fprs]  # 错误接受 负样本分错的比例
+
+    return fnrs, fprs, thresholds
+
+# Computes the minimum of the detection cost function.  The comments refer to
+# equations in Section 3 of the NIST 2016 Speaker Recognition Evaluation Plan.
+def ComputeMinDcf(fnrs, fprs, thresholds, p_target, c_miss, c_fa):
+    min_c_det = float("inf")
+    min_c_det_threshold = thresholds[0]
+    for i in range(0, len(fnrs)):
+        c_det = c_miss * fnrs[i] * p_target + c_fa * fprs[i] * (1 - p_target)
+        if c_det < min_c_det:
+            min_c_det = c_det
+            min_c_det_threshold = thresholds[i]
+
+    c_def = min(c_miss * p_target, c_fa * (1 - p_target))
+    min_dcf = min_c_det / c_def
+    return min_dcf, min_c_det_threshold