diff --git a/.gitignore b/.gitignore new file mode 100755 index 0000000..beffa3f --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +__pycache__ + +exp/ +data/ + +WavLM-Base+.pt \ No newline at end of file diff --git a/DatasetLoader.py b/DatasetLoader.py new file mode 100644 index 0000000..3b5e218 --- /dev/null +++ b/DatasetLoader.py @@ -0,0 +1,309 @@ + #! /usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +import numpy +import random +import pdb +import os +import threading +import time +import math +import glob +# import soundfile +from scipy import signal +import soundfile +from torch.utils.data import Dataset, DataLoader +import torch.distributed as dist + +def round_down(num, divisor): + return num - (num%divisor) + +def worker_init_fn(worker_id): + numpy.random.seed(numpy.random.get_state()[1][0] + worker_id) + + +def loadWAV(filename, max_frames, evalmode=True, num_eval=5): + + # Maximum audio length + max_audio = max_frames * 160 + 240 + + # Read wav file and convert to torch tensor + audio, sample_rate = soundfile.read(filename) + + + audiosize = audio.shape[0] + + if audiosize <= max_audio: + shortage = max_audio - audiosize + 1 + audio = numpy.pad(audio, (0, shortage), 'wrap') + audiosize = audio.shape[0] + + if evalmode: + startframe = numpy.linspace(0,audiosize-max_audio,num=num_eval) + else: + startframe = numpy.array([numpy.int64(random.random()*(audiosize-max_audio))]) + + feats = [] + if evalmode and max_frames == 0: + feats.append(audio) + else: + for asf in startframe: + feats.append(audio[int(asf):int(asf)+max_audio]) + + feat = numpy.stack(feats,axis=0).astype(float) + + return feat; + +class AugmentWAV(object): + + def __init__(self, musan_path, rir_path, max_frames): + + self.max_frames = max_frames + self.max_audio = max_audio = max_frames * 160 + 240 + + self.noisetypes = ['noise','speech','music'] + + self.noisesnr = {'noise':[0,15],'speech':[13,20],'music':[5,15]} + self.numnoise = {'noise':[1,1], 'speech':[3,8], 'music':[1,1] } + self.noiselist = {} + + augment_files = glob.glob(os.path.join(musan_path,'*/*/*.wav')); + + for file in augment_files: + if not file.split('/')[-3] in self.noiselist: + self.noiselist[file.split('/')[-3]] = [] + self.noiselist[file.split('/')[-3]].append(file) + + self.rir_files = glob.glob(os.path.join(rir_path,'*/*/*.wav')); + + def additive_noise(self, noisecat, audio): + + clean_db = 10 * numpy.log10(numpy.mean(audio ** 2)+1e-4) + + numnoise = self.numnoise[noisecat] + noiselist = random.sample(self.noiselist[noisecat], random.randint(numnoise[0],numnoise[1])) + + noises = [] + + for noise in noiselist: + + noiseaudio = loadWAV(noise, self.max_frames, evalmode=False) + noise_snr = random.uniform(self.noisesnr[noisecat][0],self.noisesnr[noisecat][1]) + noise_db = 10 * numpy.log10(numpy.mean(noiseaudio[0] ** 2)+1e-4) + noises.append(numpy.sqrt(10 ** ((clean_db - noise_db - noise_snr) / 10)) * noiseaudio) + + return numpy.sum(numpy.concatenate(noises,axis=0),axis=0,keepdims=True) + audio + + def reverberate(self, audio): + + rir_file = random.choice(self.rir_files) + + rir, fs = soundfile.read(rir_file) + rir = numpy.expand_dims(rir.astype(float),0) + rir = rir / numpy.sqrt(numpy.sum(rir**2)) + + return signal.convolve(audio, rir, mode='full')[:,:self.max_audio] + + +class train_dataset_loader(Dataset): + def __init__(self, train_list, augment, musan_path, rir_path, max_frames, train_path, **kwargs): + + self.augment_wav = AugmentWAV(musan_path=musan_path, rir_path=rir_path, max_frames = max_frames) + + self.train_list = train_list + self.max_frames = max_frames; + self.musan_path = musan_path + self.rir_path = rir_path + self.augment = augment + + # Read training files + with open(train_list) as dataset_file: + lines = dataset_file.readlines(); + + # Make a dictionary of ID names and ID indices + dictkeys = list(set([x.split()[0] for x in lines])) + dictkeys.sort() + dictkeys = { key : ii for ii, key in enumerate(dictkeys) } + + # Parse the training list into file names and ID indices + 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_label.append(speaker_label) + self.data_list.append(filename) + + + def __getitem__(self, indices): + + feat_clean = [] + feat = [] + + for index in indices: + try: + audio_clean = loadWAV(self.data_list[index], self.max_frames, evalmode=False) + except: + print(self.data_list[index]) + + if len(audio_clean.shape) == 3: + print(self.data_list[index]) + + if self.augment: + augtype = random.randint(0,5) + if augtype == 0: + audio = audio_clean + elif augtype == 1: + audio = self.augment_wav.reverberate(audio_clean) + elif augtype == 2: + audio = self.augment_wav.additive_noise('music',audio_clean) + elif augtype == 3: + audio = self.augment_wav.additive_noise('speech',audio_clean) + elif augtype == 4: + audio = self.augment_wav.additive_noise('noise',audio_clean) + elif augtype == 5: + audio = self.augment_wav.additive_noise('speech',audio_clean) + audio = self.augment_wav.additive_noise('music',audio_clean) + + feat_clean.append(audio_clean) + feat.append(audio) + + feat_clean = numpy.concatenate(feat_clean, axis=0) + feat = numpy.concatenate(feat, axis=0) + + return torch.FloatTensor(feat_clean), torch.FloatTensor(feat), self.data_label[index], self.data_list[index] + + def __len__(self): + return len(self.data_list) + + + +class test_dataset_loader(Dataset): + def __init__(self, test_list, test_path, eval_frames, num_eval, **kwargs): + self.max_frames = eval_frames; + self.num_eval = num_eval + self.test_path = test_path + self.test_list = test_list + + def __getitem__(self, index): + # print(self.test_list[index]) + audio = loadWAV(os.path.join(self.test_path,self.test_list[index]), self.max_frames, evalmode=True, num_eval=self.num_eval) + + audio2 = loadWAV(os.path.join(self.test_path,self.test_list[index]), 0, evalmode=True, num_eval=self.num_eval) + + return torch.FloatTensor(audio), torch.FloatTensor(audio2), self.test_list[index] + # return torch.FloatTensor(audio2), self.test_list[index] + + def __len__(self): + return len(self.test_list) + + +class train_dataset_sampler(torch.utils.data.Sampler): + def __init__(self, data_source, nPerSpeaker, max_seg_per_spk, batch_size, distributed, seed, **kwargs): + + self.data_label = data_source.data_label; + self.nPerSpeaker = nPerSpeaker; + self.max_seg_per_spk = max_seg_per_spk; + self.batch_size = batch_size; + self.epoch = 0; + self.seed = seed; + self.distributed = distributed; + + def __iter__(self): + + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.data_label), generator=g).tolist() + + data_dict = {} + + # Sort into dictionary of file indices for each ID + for index in indices: + speaker_label = self.data_label[index] + if not (speaker_label in data_dict): + data_dict[speaker_label] = []; + data_dict[speaker_label].append(index); + + + ## Group file indices for each class + dictkeys = list(data_dict.keys()); + dictkeys.sort() + + lol = lambda lst, sz: [lst[i:i+sz] for i in range(0, len(lst), sz)] + + flattened_list = [] + flattened_label = [] + + for findex, key in enumerate(dictkeys): + data = data_dict[key] + numSeg = round_down(min(len(data),self.max_seg_per_spk),self.nPerSpeaker) + + rp = lol(numpy.arange(numSeg),self.nPerSpeaker) + flattened_label.extend([findex] * (len(rp))) + for indices in rp: + flattened_list.append([data[i] for i in indices]) + + ## Mix data in random order + mixid = torch.randperm(len(flattened_label), generator=g).tolist() + mixlabel = [] + mixmap = [] + + ## Prevent two pairs of the same speaker in the same batch + for ii in mixid: + startbatch = round_down(len(mixlabel), self.batch_size) + if flattened_label[ii] not in mixlabel[startbatch:]: + mixlabel.append(flattened_label[ii]) + mixmap.append(ii) + + mixed_list = [flattened_list[i] for i in mixmap] + + ## Divide data to each GPU + if self.distributed: + total_size = round_down(len(mixed_list), self.batch_size * dist.get_world_size()) + start_index = int ( ( dist.get_rank() ) / dist.get_world_size() * total_size ) + end_index = int ( ( dist.get_rank() + 1 ) / dist.get_world_size() * total_size ) + self.num_samples = end_index - start_index + return iter(mixed_list[start_index:end_index]) + else: + total_size = round_down(len(mixed_list), self.batch_size) + self.num_samples = total_size + return iter(mixed_list[:total_size]) + + + def __len__(self) -> int: + return self.num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + +if __name__ == '__main__': + train_dataset = train_dataset_loader(train_list='/mnt/proj3/open-24-5/pengjy_new/WavLM_Adapter/CNCeleb_lst/CNCeleb_trainlist_200spk.txt', + augment=False, + musan_path='/mnt/proj3/open-24-5/pengjy_new/musan_split/', + rir_path='/mnt/proj3/open-24-5/plchot/data_augment/16kHz/simulated_rirs/', + max_frames=300, + train_path='/mnt/proj3/open-24-5/pengjy_new/Data/CN-Celeb_flac/data', + ) + + train_sampler = train_dataset_sampler(train_dataset, nPerSpeaker=1, max_seg_per_spk=500, batch_size=100, distributed=False,seed=120) + # train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) + + train_loader = torch.utils.data.DataLoader( + train_dataset, + batch_size=100, + num_workers=10, + sampler=train_sampler, + pin_memory=True, + drop_last=True, + ) + for data, data_label in train_loader: + print(data.shape) + data = data.transpose(1,0) + print(data.shape) + quit() \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..d8f84cb --- /dev/null +++ b/README.md @@ -0,0 +1,96 @@ +# wavlm_ssl_sv + +This repository contains the source code of the article **Towards Supervised Performance on Speaker Verification with Self-Supervised Learning by Leveraging Large-Scale ASR Models** (INTERSPEECH 2024) [[arXiv]](https://arxiv.org/pdf/2406.02285). + +The proposed framework fine-tunes a pre-trained **WavLM** using pseudo-labels, generated through **Self-Supervised Learning** (SSL), for **Speaker Verification** (SV). Initial pseudo-labels are derived from an SSL DINO-based model and are iteratively refined by clustering the model embeddings. + +

+ +

+ +Our method achieves **0.99% EER on VoxCeleb1-O**, establishing the new SOTA on Speaker Verification with SSL. + +*Please refer to the article for more details on the implementation and a comparative study with other works.* + +--- + +## Usage + +### Installation + +- Install dependencies with `pip install -r requirements.txt`. +- Prepare data for VoxCeleb, MUSAN, and RIR datasets following [voxceleb_trainer](https://github.com/clovaai/voxceleb_trainer#data-preparation). +- Download [WavLM-Base+ model](https://github.com/microsoft/unilm/tree/master/wavlm) and place `WavLM-Base+.pt` at the root folder. + +### Training + +#### Step 1: Extract DINO speaker embeddings + +The code to train the DINO model is not currently provided. We recommend using [sslsv](https://github.com/theolepage/sslsv) or [3D-Speaker](https://github.com/modelscope/3D-Speaker) to extract initial speaker embeddings. + +Alternatively, you can directly download the DINO embeddings we used for our system: [dino_vox2_embeddings.pt](https://drive.google.com/file/d/1YnxrMIgrr6NQgZ3Hv2_5YdP5W8xfdyLH/view?usp=sharing). + +*Note: the embeddings file must be a `Dict[str, torch.Tensor]` representing all VoxCeleb2 samples with the following format for keys: `id00012/21Uxsk56VDQ/00001.wav`.* + +#### Step 2: Generate pseudo-labels + +```bash +python pseudo_labeling.py PATH_TO_EMBEDDINGS_FILE PATH_TO_PL_FILE +``` + +#### Step 3: Fine-tune WavLM MHFA + +```bash +python trainSpeakerNet.py --config configs/wavlm_mhfa_dlg_lc.yaml --train_list PATH_TO_PL_FILE --distributed +``` + +#### Iterative process + +1. Extract embeddings from the WavLM MHFA model: + `python trainSpeakerNet_Eval.py --config configs/wavlm_mhfa_dlg_lc.yaml --generate_embeddings --embeddings_path PATH_TO_EMBEDDINGS_FILE`. + +2. Repeat steps 2 and 3. *Make sure to change `save_path` in the config to avoid overwriting the existing model.* + +#### Step 4: Large-Margin Fine-Tuning + +1. Copy the latest model checkpoint to `exp/wavlm_mhfa_dlg_lc_lmft/model` to resume training. + +2. Start training: `python trainSpeakerNet.py --config configs/wavlm_mhfa_dlg_lc_lmft.yaml --train_list PATH_TO_PL_FILE --distributed`. + +### Evaluation + +```bash +python trainSpeakerNet_Eval.py --config configs/wavlm_mhfa_dlg_lc_lmft.yaml --eval +``` + +### Model weights + +The checkpoint of our best model reaching 0.99% EER on VoxCeleb1-O is available for download: [`wavlm_mhfa_dlg_lc_lmft`](https://drive.google.com/drive/folders/1ygZPvdGwepWDDfIQp6aPRktt2QxLt6cE?usp=drive_link). + +--- + +## Acknowledgements + +This repository contains third-party components and code adapted from other open-source projects, including: [SLT22_MultiHead-Factorized-Attentive-Pooling](https://github.com/JunyiPeng00/SLT22_MultiHead-Factorized-Attentive-Pooling) and [Loss-Gated-Learning](https://github.com/TaoRuijie/Loss-Gated-Learning). + +--- + +## Citation + +If you use this project, please consider starring this repository on GitHub and citing the following paper. + +```BibTeX +@InProceedings{miara2024WavLMSSLSV, + author = {Miara, Victor and Lepage, Théo and Dehak, Réda}, + booktitle = {INTERSPEECH}, + title = {Towards Supervised Performance on Speaker Verification with Self-Supervised Learning by Leveraging Large-Scale ASR Models}, + year = {2024}, + url = {https://arxiv.org/abs/2406.02285}, +} +``` + +--- + +## License + +This project is released under the [MIT License](https://github.com/theolepage/wavlm_ssl_sv/blob/main/LICENSE.md). diff --git a/SpeakerNet.py b/SpeakerNet.py new file mode 100644 index 0000000..d5c10d3 --- /dev/null +++ b/SpeakerNet.py @@ -0,0 +1,372 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy, math, pdb, sys, random +import time, os, itertools, shutil, importlib +from tuneThreshold import tuneThresholdfromScore +from DatasetLoader import test_dataset_loader, loadWAV +import pickle +import numpy as np +import time +from tqdm import tqdm +import soundfile + + +class WrappedModel(nn.Module): + + ## The purpose of this wrapper is to make the model structure consistent between single and multi-GPU + + def __init__(self, model): + super(WrappedModel, self).__init__() + self.module = model + + def forward(self, x, x_clean=None, label=None,l2_reg_dict=None, epoch=-1): + return self.module(x, x_clean, label, epoch=epoch) + + +class SpeakerNet(nn.Module): + + def __init__(self, model, optimizer, trainfunc, nPerSpeaker, **kwargs): + super(SpeakerNet, self).__init__() + + SpeakerNetModel = importlib.import_module('models.'+model).__getattribute__('MainModel') + self.__S__ = SpeakerNetModel(**kwargs); + + LossFunction = importlib.import_module('loss.'+trainfunc).__getattribute__('LossFunction') + self.__L__ = LossFunction(**kwargs); + + self.nPerSpeaker = nPerSpeaker + self.weight_finetuning_reg = kwargs['weight_finetuning_reg'] + + + def forward(self, data, data_clean=None, label=None, l2_reg_dict=None, epoch=-1): + if label is None: + data_reshape = data[0].cuda() + outp = self.__S__.forward([data_reshape, data[1]]) + return outp + elif len(data) == 3 and data[2] == "gen_ps": + data_reshape = data[0].reshape(-1,data[0].size()[-1]).cuda() + outp = self.__S__.forward([data_reshape, data[1]]) + pseudo_labels = self.__L__.get_pseudo_labels(outp, label) + return pseudo_labels + else: + data_reshape = data[0].reshape(-1,data[0].size()[-1]).cuda() + data_clean_reshape = data_clean.reshape(-1,data_clean.size()[-1]).cuda() + outp = self.__S__.forward([data_reshape, data[1]]) + outp_clean = self.__S__.forward([data_clean_reshape, data[1]]) + nloss, prec1, ce = self.__L__.forward(outp, outp_clean, label, epoch) + + if l2_reg_dict is not None: + Learned_dict = l2_reg_dict + l2_reg = 0 + for name,param in self.__S__.model.named_parameters(): + if name in Learned_dict: + l2_reg = l2_reg + torch.norm(param-Learned_dict[name].cuda(),2) + tloss = nloss/nloss.detach() + self.weight_finetuning_reg*l2_reg/(l2_reg.detach()+1e-5) + else: + tloss = nloss + print("Without L2 Reg") + + return tloss, prec1, nloss, ce + + + + +class ModelTrainer(object): + + def __init__(self, speaker_model, optimizer, scheduler, gpu, mixedprec, **kwargs): + + self.__model__ = speaker_model + + WavLM_params = list(map(id, self.__model__.module.__S__.model.parameters())) + Backend_params = filter(lambda p: id(p) not in WavLM_params, self.__model__.module.parameters()) + self.path = kwargs['pretrained_model_path'] + + Optimizer = importlib.import_module('optimizer.'+optimizer).__getattribute__('Optimizer') + + # Define the initial param groups + param_groups = [{'params': Backend_params, 'lr': kwargs['LR_MHFA']}] + + # Extract the encoder layers + encoder_layers = self.__model__.module.__S__.model.encoder.layers + + # Iterate over the encoder layers to create param groups + for i in range(12): # Assuming 12 layers from 0 to 11 (for BASE model, when it comes to LARGE model, 12->24) + lr = kwargs['LR_Transformer'] * (kwargs['LLRD_factor'] ** i) + param_groups.append({'params': encoder_layers[i].parameters(), 'lr': lr}) + + # Initialize the optimizer with these param groups + self.__optimizer__ = Optimizer(param_groups, **kwargs) + + # self.__optimizer__ = Optimizer(self.__model__.parameters(), **kwargs) + # print('scheduler.'+scheduler) + Scheduler = importlib.import_module('scheduler.'+scheduler).__getattribute__('Scheduler') + # print(kwargs) + try: + self.__scheduler__, self.lr_step = Scheduler(self.__optimizer__, **kwargs) + except: + self.__scheduler__, self.lr_step = Scheduler(self.__optimizer__, lr_decay=0.9, **kwargs) + + # self.scaler = GradScaler() + + self.gpu = gpu + + self.mixedprec = mixedprec + print("Mix prec: %s"%(self.mixedprec)) + + assert self.lr_step in ['epoch', 'iteration'] + + # ## ===== ===== ===== ===== ===== ===== ===== ===== + # ## Train network + # ## ===== ===== ===== ===== ===== ===== ===== ===== + + def update_lgl_threshold(self, lgl_threshold): + self.__model__.module.__L__.lgl_threshold = lgl_threshold + + # """ + def train_network(self, loader, loss_vals_path, epoch, verbose): + if torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + unique_loss_vals_path = f"{loss_vals_path.split('.')[0]}_rank{rank}.txt" + else: + unique_loss_vals_path = loss_vals_path + + self.__model__.train(); + + stepsize = loader.batch_size; + + counter = 0; + index = 0; + loss = 0; + top1 = 0 # EER or accuracy + + tstart = time.time() + Learned_dict = {} + checkpoint = torch.load(self.path) + for name, param in checkpoint['model'].items(): + if 'w2v_encoder.w2v_model.' in name: + newname = name.replace('w2v_encoder.w2v_model.', '') + else: + newname = name + Learned_dict[newname] = param; + + # for data_clean, data, data_label, data_path in loader: + # telapsed = time.time() - tstart + # tstart = time.time() + # counter += 1; + # index += stepsize + # sys.stdout.write("\rProcessing (%d) "%(index)); + # sys.stdout.write("Loss %f TEER/TAcc %2.3f%% - %.2f Hz "%(loss/counter, top1/counter, stepsize/telapsed)); + # if counter % 100 == 0: + # sys.stdout.flush() + + with open(unique_loss_vals_path, 'w') as loss_vals_file: + for data_clean, data, data_label, data_path in loader: + data_clean = data_clean.transpose(1,0) + data = data.transpose(1,0) + self.__model__.zero_grad() + label = torch.LongTensor(data_label).cuda() + + nloss, prec1, spkloss, ce = self.__model__([data,"train"], data_clean, label, Learned_dict, epoch=epoch) + + for ce_val, path in zip(ce.detach().cpu().numpy(), data_path): + loss_vals_file.write(f'{ce_val} {"/".join(path.split("/")[5:])}\n') + + nloss.backward() + + self.__optimizer__.step(); + + loss += spkloss.detach().cpu() + top1 += prec1.detach().cpu() + + + counter += 1; + index += stepsize; + + + + telapsed = time.time() - tstart + tstart = time.time() + + if verbose: + sys.stdout.write("\rProcessing (%d) "%(index)); + sys.stdout.write("Loss %f TEER/TAcc %2.3f%% - %.2f Hz "%(loss/counter, top1/counter, stepsize/telapsed)); + sys.stdout.flush(); + + if self.lr_step == 'iteration': self.__scheduler__.step() + + if self.lr_step == 'epoch': self.__scheduler__.step() + + sys.stdout.write("\n"); + return (loss/counter, top1/counter); + # """ + + ## ===== ===== ===== ===== ===== ===== ===== ===== + ## Evaluate from list + ## ===== ===== ===== ===== ===== ===== ===== ===== + + def evaluateFromList(self, test_list, test_path, nDataLoaderThread, print_interval=10, num_eval=15, **kwargs): + + self.__model__.eval(); + + lines = [] + files = [] + feats = {} + tstart = time.time() + + ## Read all lines + with open(test_list) as f: + lines = f.readlines() + + ## Get a list of unique file names + files = sum([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, num_eval=num_eval, **kwargs) + test_loader = torch.utils.data.DataLoader( + test_dataset, + batch_size=1, + shuffle=False, + num_workers=nDataLoaderThread, + drop_last=False, + ) + ref_feat_list = [] + ref_feat_2_list = [] + max_len = 0 + forward = 0 + ## Extract features for every image + for idx, data in enumerate(test_loader): + + + inp1 = data[0][0].cuda() + inp2 = data[1][0].cuda() + telapsed_2 = time.time() + b,utt_l = inp2.shape + if utt_l > max_len: + max_len = utt_l + ref_feat = self.__model__([inp1, "test"]).cuda() + ref_feat = ref_feat.detach().cpu() + ref_feat_2 = self.__model__([inp2[:,:700000], "test"]).cuda() # The reason why here is set to 700000 is due to GPU memory size. + ref_feat_2 = ref_feat_2.detach().cpu() + + feats[data[2][0]] = [ref_feat,ref_feat_2] + + ref_feat_list.extend(ref_feat.numpy()) + ref_feat_2_list.extend(ref_feat_2.numpy()) + + telapsed = time.time() - tstart + forward = forward + time.time() - telapsed_2 + + if idx % print_interval == 0: + sys.stdout.write("\rReading %d of %d: %.2f Hz, forward speed: %.2f Hz, embedding size %d, max_len %d"%(idx,len(setfiles),idx/telapsed,idx/forward, ref_feat.size()[-1],max_len)); + + print('') + all_scores = []; + all_labels = []; + all_trials = []; + all_scores_1 = []; + all_scores_2 = []; + + tstart = time.time() + + ref_feat_list = numpy.array(ref_feat_list) + ref_feat_2_list = numpy.array(ref_feat_2_list) + + ref_feat_list_mean = 0 + ref_feat_2_list_mean = 0 + + + ## Read files and compute all scores + for idx, line in enumerate(lines): + + data = line.split(); + + ## Append random label if missing + if len(data) == 2: data = [random.randint(0,1)] + data + + ref_feat,ref_feat_2 = feats[data[1]] + com_feat,com_feat_2 = feats[data[2]] + + # if self.__model__.module.__L__.test_normalize: + ref_feat = F.normalize(ref_feat-ref_feat_list_mean, p=2, dim=1) # B, D + com_feat = F.normalize(com_feat-ref_feat_list_mean, p=2, dim=1) + ref_feat_2 = F.normalize(ref_feat_2-ref_feat_2_list_mean, p=2, dim=1) # B, D + com_feat_2 = F.normalize(com_feat_2-ref_feat_2_list_mean, p=2, dim=1) + + score_1 = torch.mean(torch.matmul(ref_feat, com_feat.T)) # higher is positive + score_2 = torch.mean(torch.matmul(ref_feat_2, com_feat_2.T)) + score = (score_1 + score_2) / 2 + score = score.detach().cpu().numpy() + + all_scores.append(score); + all_scores_1.append(score_1); + all_scores_2.append(score_2); + + all_labels.append(int(data[0])); + all_trials.append(data[1]+" "+data[2]) + + if idx % (10*print_interval) == 0: + telapsed = time.time() - tstart + sys.stdout.write("\rComputing %d of %d: %.2f Hz"%(idx,len(lines),idx/telapsed)); + sys.stdout.flush(); + + print('') + + return (all_scores, all_labels, all_trials,all_scores_1,all_scores_2); + + def generate_embeddings(self, wav_files, output, device): + res = {} + + for file in tqdm(wav_files): + wav, sr = soundfile.read(file) + wav = torch.from_numpy(wav).float().to(device) + + with torch.no_grad(): + embedding = self.__model__([wav.unsqueeze(0), "test"]).detach().cpu() + + key = '/'.join(file.split('/')[-3:]) + res[key] = embedding + + torch.save(res, output) + + def saveParameters(self, path): + torch.save(self.__model__.module.state_dict(), path); + + + ## ===== ===== ===== ===== ===== ===== ===== ===== + ## Load parameters + ## ===== ===== ===== ===== ===== ===== ===== ===== + + def loadParameters(self, path): + + self_state = self.__model__.module.state_dict(); + loaded_state = torch.load(path, map_location="cuda:%d"%self.gpu); + # loaded_state = torch.load(path, map_location="cpu"); + + + + for name, param in loaded_state.items(): + origname = name; + + if name not in self_state: + name = name.replace("module.", ""); + + if name not in self_state: + print("%s is not in the model."%origname); + continue; + + if self_state[name].size() != loaded_state[origname].size(): + print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, self_state[name].size(), loaded_state[origname].size())); + continue; + + self_state[name].copy_(param); + + + + + diff --git a/configs/wavlm_mhfa_dlg_lc.yaml b/configs/wavlm_mhfa_dlg_lc.yaml new file mode 100644 index 0000000..d3a3d9e --- /dev/null +++ b/configs/wavlm_mhfa_dlg_lc.yaml @@ -0,0 +1,33 @@ +max_frames: 300 +max_epoch: 15 +batch_size: 120 +margin: 0.2 + +eval_frames: 400 +augment: True + +## Training details +trainfunc: aamsoftmax + +scale: 30 + +lr_decay: 0.95 + +pretrained_model_path: WavLM-Base+.pt +weight_finetuning_reg: 0.01 +LLRD_factor: 1.0 +LR_Transformer: 2e-5 +LR_MHFA: 5e-3 + +## Loss functions +nClasses: 7500 + +## Load and save +save_path: exp/wavlm_mhfa_dlg_lc +# save_path: exp/wavlm_mhfa_dlg_lc_iter2 + +## Model definition +model: Baseline.Spk_Encoder + +nOut: 256 +port: 6754 diff --git a/configs/wavlm_mhfa_dlg_lc_lmft.yaml b/configs/wavlm_mhfa_dlg_lc_lmft.yaml new file mode 100644 index 0000000..d182dee --- /dev/null +++ b/configs/wavlm_mhfa_dlg_lc_lmft.yaml @@ -0,0 +1,33 @@ +max_frames: 600 +max_epoch: 5 +batch_size: 50 +margin: 0.5 + +eval_frames: 400 +augment: True + +## Training details +trainfunc: aamsoftmax + +scale: 30 + +lr: 5e-4 +lr_decay: 0.95 + +pretrained_model_path: WavLM-Base+.pt +weight_finetuning_reg: 0.01 +LLRD_factor: 1.0 +LR_Transformer: 2e-5 +LR_MHFA: 5e-3 + +## Loss functions +nClasses: 7500 + +## Load and save +save_path: exp/wavlm_mhfa_dlg_lc_lmft + +## Model definition +model: Baseline.Spk_Encoder + +nOut: 256 +port: 6754 diff --git a/loss/aamsoftmax.py b/loss/aamsoftmax.py new file mode 100644 index 0000000..27065b0 --- /dev/null +++ b/loss/aamsoftmax.py @@ -0,0 +1,140 @@ +#! /usr/bin/python +# -*- encoding: utf-8 -*- +# Adapted from https://github.com/wujiyang/Face_Pytorch (Apache License) + +import torch +import torch.nn as nn +import torch.nn.functional as F +import time, pdb, numpy, math +from utils import accuracy +import numpy as np + +class LossFunction(nn.Module): + def __init__(self, nOut, nClasses, margin=0.3, scale=15, easy_margin=False, **kwargs): + super(LossFunction, self).__init__() + + self.test_normalize = True + + self.m = margin + self.s = scale + self.in_feats = nOut + self.weight = torch.nn.Parameter(torch.FloatTensor(nClasses, nOut), requires_grad=True) + # self.ce = nn.CrossEntropyLoss() + self.ce = nn.CrossEntropyLoss(reduction='none') # return loss per sample + nn.init.xavier_normal_(self.weight, gain=1) + + self.easy_margin = easy_margin + self.cos_m = math.cos(self.m) + self.sin_m = math.sin(self.m) + + # make the function cos(theta+m) monotonic decreasing while theta in [0°,180°] + self.th = math.cos(math.pi - self.m) + self.mm = math.sin(math.pi - self.m) * self.m + + self.lgl_threshold = 1e6 + self.lc_threshold = 0.5 + + print('Initialised AAMSoftmax margin %.3f scale %.3f'%(self.m,self.s)) + + def _forward(self, x, label): + # cos(theta) + cosine = F.linear(F.normalize(x), F.normalize(self.weight)) + # cos(theta + m) + sine = torch.sqrt((1.0 - torch.mul(cosine, cosine)).clamp(0, 1)) + phi = cosine * self.cos_m - sine * self.sin_m + + if self.easy_margin: + phi = torch.where(cosine > 0, phi, cosine) + else: + phi = torch.where((cosine - self.th) > 0, phi, cosine - self.mm) + + #one_hot = torch.zeros(cosine.size(), device='cuda' if torch.cuda.is_available() else 'cpu') + one_hot = torch.zeros_like(cosine) + one_hot.scatter_(1, label.view(-1, 1), 1) + output = (one_hot * phi) + ((1.0 - one_hot) * cosine) + output = output * self.s + + return output + + def _forward_softmax_sharpened(self, x, e=0.1): + # regular softmax + output = F.linear(x, self.weight) + probas = F.softmax(output / e, dim=1) + return probas + + def forward(self, x, x_clean, label=None, epoch=-1): + assert x.size()[0] == label.size()[0] + assert x.size()[1] == self.in_feats + + output = self._forward(x, label) + output_clean = self._forward_softmax_sharpened(x_clean) + + ce = self.ce(output, label) + + # No LGL + # prec1 = accuracy(output.detach(), label.detach(), topk=(1,))[0] + # return ce, prec1, None + + mask = (torch.log(ce) <= self.lgl_threshold).detach() + + if epoch <= 8: + # LGL only + nselect = torch.clamp(sum(mask), min=1).item() + loss = torch.sum(ce * mask, dim=-1) / nselect + prec1 = accuracy(output.detach(), label * mask.detach(), topk=(1,))[0] + return loss, prec1, ce + + # LGL + LC + + label_LC = output_clean.argmax(dim=1) + + max_vals = torch.gather(output_clean, 1, label_LC.unsqueeze(1)).squeeze(1) + mask_LC = (max_vals > self.lc_threshold).detach() + + ce_LC = self.ce(output, label_LC) + + mask_LGL_LC = ~mask & mask_LC + loss = torch.mean(ce * mask + ce_LC * mask_LGL_LC, dim=-1) + prec1 = accuracy(output.detach(), label * mask.detach() + label_LC * mask_LGL_LC.detach(), topk=(1,))[0] + + return loss, prec1, ce + + def get_pseudo_labels(self, x, label): + output = self._forward_softmax_sharpened(x) + return output.argmax(dim=1) + +""" + def forward(self, x, x_clean, label=None): + + assert x.size()[0] == label.size()[0] + assert x.size()[1] == self.in_feats + + P_aam = self._forward(x, label) + + P_softmax = self._forward_softmax_sharpened(x) + P_clean_softmax = self._forward_softmax_sharpened(x_clean) + + ce = self.ce(P_aam, label) + + # No LGL + # prec1 = accuracy(output.detach(), label.detach(), topk=(1,))[0] + # return ce, prec1, None + + mask = (torch.log(ce) <= self.lgl_threshold).detach() + + # LGL only + # nselect = torch.clamp(sum(mask), min=1).item() + # loss = torch.sum(ce * mask, dim=-1) / nselect + # prec1 = accuracy(output.detach(), label * mask.detach(), topk=(1,))[0] + # return loss, prec1, ce + + # LGL + LC + label_LC = P_clean_softmax.argmax(dim=1) + ce_LC = self.ce(P_softmax, label_LC) + + inverted_mask = ~mask + loss = torch.mean(ce * mask + ce_LC * inverted_mask, dim=-1) + prec1 = accuracy(P_softmax.detach(), label * mask.detach() + label_LC * inverted_mask.detach(), topk=(1,))[0] + + return loss, prec1, ce +""" \ No newline at end of file diff --git a/models/Baseline/Spk_Encoder.py b/models/Baseline/Spk_Encoder.py new file mode 100644 index 0000000..dc80870 --- /dev/null +++ b/models/Baseline/Spk_Encoder.py @@ -0,0 +1,112 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import LayerNorm +from .WavLM import * + +import torch +import torch.nn as nn + +class MHFA(nn.Module): + def __init__(self, head_nb=8, inputs_dim=768, compression_dim=128, outputs_dim=256): + super(MHFA, self).__init__() + + # Define learnable weights for key and value computations across layers + self.weights_k = nn.Parameter(data=torch.ones(13), requires_grad=True) + self.weights_v = nn.Parameter(data=torch.ones(13), requires_grad=True) + + # Initialize given parameters + self.head_nb = head_nb + self.ins_dim = inputs_dim + self.cmp_dim = compression_dim + self.ous_dim = outputs_dim + + # Define compression linear layers for keys and values + self.cmp_linear_k = nn.Linear(self.ins_dim, self.cmp_dim) + self.cmp_linear_v = nn.Linear(self.ins_dim, self.cmp_dim) + + # Define linear layer to compute multi-head attention weights + self.att_head = nn.Linear(self.cmp_dim, self.head_nb) + + # Define a fully connected layer for final output + self.pooling_fc = nn.Linear(self.head_nb * self.cmp_dim, self.ous_dim) + + def forward(self, x): + # Input x has shape: [Batch, Dim, Frame_len, Nb_Layer] + + # Compute the key by taking a weighted sum of input across layers + k = torch.sum(x.mul(nn.functional.softmax(self.weights_k, dim=-1)), dim=-1).transpose(1, 2) + + # Compute the value in a similar fashion + v = torch.sum(x.mul(nn.functional.softmax(self.weights_v, dim=-1)), dim=-1).transpose(1, 2) + + # Pass the keys and values through compression linear layers + k = self.cmp_linear_k(k) + v = self.cmp_linear_v(v) + + # Compute attention weights using compressed keys + att_k = self.att_head(k) + + # Adjust dimensions for computing attention output + v = v.unsqueeze(-2) + + # Compute attention output by taking weighted sum of values using softmaxed attention weights + pooling_outs = torch.sum(v.mul(nn.functional.softmax(att_k, dim=1).unsqueeze(-1)), dim=1) + + # Reshape the tensor before passing through the fully connected layer + b, h, f = pooling_outs.shape + pooling_outs = pooling_outs.reshape(b, -1) + + # Pass through fully connected layer to get the final output + outs = self.pooling_fc(pooling_outs) + + return outs + + +class spk_extractor(nn.Module): + def __init__(self,**kwargs): + super(spk_extractor, self).__init__() + # checkpoint = torch.load('/mnt/proj3/open-24-5/pengjy_new/WavLM/Pretrained_model/WavLM-Base+.pt') + print("Pre-trained Model: {}".format(kwargs['pretrained_model_path'])) + checkpoint = torch.load(kwargs['pretrained_model_path']) + cfg = WavLMConfig(checkpoint['cfg']) + self.model = WavLM(cfg) + self.loadParameters(checkpoint['model']) + self.backend = MHFA(head_nb=64) + + + def forward(self,wav_and_flag): + + x = wav_and_flag[0] + + cnn_outs, layer_results = self.model.extract_features(x, output_layer=13) + layer_reps = [x.transpose(0, 1) for x, _ in layer_results] + x = torch.stack(layer_reps).transpose(0,-1).transpose(0,1) + + out = self.backend(x) + return out + + def loadParameters(self, param): + + self_state = self.model.state_dict(); + loaded_state = param + + for name, param in loaded_state.items(): + origname = name; + + + if name not in self_state: + # print("%s is not in the model."%origname); + continue; + + if self_state[name].size() != loaded_state[origname].size(): + print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, self_state[name].size(), loaded_state[origname].size())); + continue; + + self_state[name].copy_(param); + + +def MainModel(**kwargs): + model = spk_extractor(**kwargs) + return model diff --git a/models/Baseline/WavLM.py b/models/Baseline/WavLM.py new file mode 100644 index 0000000..dfa2147 --- /dev/null +++ b/models/Baseline/WavLM.py @@ -0,0 +1,749 @@ +# -------------------------------------------------------- +# WavLM: Large-Scale Self-Supervised Pre-training for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf) +# Github source: https://github.com/microsoft/unilm/tree/master/wavlm +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Based on fairseq code bases +# https://github.com/pytorch/fairseq +# -------------------------------------------------------- + +import math +import logging +from typing import List, Optional, Tuple + +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import LayerNorm +from .modules import ( + Fp32GroupNorm, + Fp32LayerNorm, + GradMultiply, + MultiheadAttention, + SamePad, + init_bert_params, + get_activation_fn, + TransposeLast, + GLU_Linear, +) + +logger = logging.getLogger(__name__) + + +def compute_mask_indices( + shape: Tuple[int, int], + padding_mask: Optional[torch.Tensor], + mask_prob: float, + mask_length: int, + mask_type: str = "static", + mask_other: float = 0.0, + min_masks: int = 0, + no_overlap: bool = False, + min_space: int = 0, +) -> np.ndarray: + """ + Computes random mask spans for a given shape + + Args: + shape: the the shape for which to compute masks. + should be of size 2 where first element is batch size and 2nd is timesteps + padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements + mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by + number of timesteps divided by length of mask span to mask approximately this percentage of all elements. + however due to overlaps, the actual number will be smaller (unless no_overlap is True) + mask_type: how to compute mask lengths + static = fixed size + uniform = sample from uniform distribution [mask_other, mask_length*2] + normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element + poisson = sample from possion distribution with lambda = mask length + min_masks: minimum number of masked spans + no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping + min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans + """ + + bsz, all_sz = shape + mask = np.full((bsz, all_sz), False) + + all_num_mask = int( + # add a random number for probabilistic rounding + mask_prob * all_sz / float(mask_length) + + np.random.rand() + ) + + all_num_mask = max(min_masks, all_num_mask) + + mask_idcs = [] + for i in range(bsz): + if padding_mask is not None: + sz = all_sz - padding_mask[i].long().sum().item() + num_mask = int( + # add a random number for probabilistic rounding + mask_prob * sz / float(mask_length) + + np.random.rand() + ) + num_mask = max(min_masks, num_mask) + else: + sz = all_sz + num_mask = all_num_mask + + if mask_type == "static": + lengths = np.full(num_mask, mask_length) + elif mask_type == "uniform": + lengths = np.random.randint(mask_other, mask_length * 2 + 1, size=num_mask) + elif mask_type == "normal": + lengths = np.random.normal(mask_length, mask_other, size=num_mask) + lengths = [max(1, int(round(x))) for x in lengths] + elif mask_type == "poisson": + lengths = np.random.poisson(mask_length, size=num_mask) + lengths = [int(round(x)) for x in lengths] + else: + raise Exception("unknown mask selection " + mask_type) + + if sum(lengths) == 0: + lengths[0] = min(mask_length, sz - 1) + + if no_overlap: + mask_idc = [] + + def arrange(s, e, length, keep_length): + span_start = np.random.randint(s, e - length) + mask_idc.extend(span_start + i for i in range(length)) + + new_parts = [] + if span_start - s - min_space >= keep_length: + new_parts.append((s, span_start - min_space + 1)) + if e - span_start - keep_length - min_space > keep_length: + new_parts.append((span_start + length + min_space, e)) + return new_parts + + parts = [(0, sz)] + min_length = min(lengths) + for length in sorted(lengths, reverse=True): + lens = np.fromiter( + (e - s if e - s >= length + min_space else 0 for s, e in parts), + np.int, + ) + l_sum = np.sum(lens) + if l_sum == 0: + break + probs = lens / np.sum(lens) + c = np.random.choice(len(parts), p=probs) + s, e = parts.pop(c) + parts.extend(arrange(s, e, length, min_length)) + mask_idc = np.asarray(mask_idc) + else: + min_len = min(lengths) + if sz - min_len <= num_mask: + min_len = sz - num_mask - 1 + + mask_idc = np.random.choice(sz - min_len, num_mask, replace=False) + + mask_idc = np.asarray( + [ + mask_idc[j] + offset + for j in range(len(mask_idc)) + for offset in range(lengths[j]) + ] + ) + + mask_idcs.append(np.unique(mask_idc[mask_idc < sz])) + + min_len = min([len(m) for m in mask_idcs]) + for i, mask_idc in enumerate(mask_idcs): + if len(mask_idc) > min_len: + mask_idc = np.random.choice(mask_idc, min_len, replace=False) + mask[i, mask_idc] = True + + return mask + + +class WavLMConfig: + def __init__(self, cfg=None): + self.extractor_mode: str = "default" # mode for feature extractor. default has a single group norm with d groups in the first conv block, whereas layer_norm has layer norms in every block (meant to use with normalize=True) + self.encoder_layers: int = 12 # num encoder layers in the transformer + + self.encoder_embed_dim: int = 768 # encoder embedding dimension + self.encoder_ffn_embed_dim: int = 3072 # encoder embedding dimension for FFN + self.encoder_attention_heads: int = 12 # num encoder attention heads + self.activation_fn: str = "gelu" # activation function to use + + self.layer_norm_first: bool = False # apply layernorm first in the transformer + self.conv_feature_layers: str = "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2" # string describing convolutional feature extraction layers in form of a python list that contains [(dim, kernel_size, stride), ...] + self.conv_bias: bool = False # include bias in conv encoder + self.feature_grad_mult: float = 1.0 # multiply feature extractor var grads by this + + self.normalize: bool = False # normalize input to have 0 mean and unit variance during training + + # dropouts + self.dropout: float = 0.1 # dropout probability for the transformer + self.attention_dropout: float = 0.1 # dropout probability for attention weights + self.activation_dropout: float = 0.0 # dropout probability after activation in FFN + self.encoder_layerdrop: float = 0.0 # probability of dropping a tarnsformer layer + self.dropout_input: float = 0.0 # dropout to apply to the input (after feat extr) + self.dropout_features: float = 0.0 # dropout to apply to the features (after feat extr) + + # masking + self.mask_length: int = 10 # mask length + self.mask_prob: float = 0.65 # probability of replacing a token with mask + self.mask_selection: str = "static" # how to choose mask length + self.mask_other: float = 0 # secondary mask argument (used for more complex distributions), see help in compute_mask_indicesh + self.no_mask_overlap: bool = False # whether to allow masks to overlap + self.mask_min_space: int = 1 # min space between spans (if no overlap is enabled) + + # channel masking + self.mask_channel_length: int = 10 # length of the mask for features (channels) + self.mask_channel_prob: float = 0.0 # probability of replacing a feature with 0 + self.mask_channel_selection: str = "static" # how to choose mask length for channel masking + self.mask_channel_other: float = 0 # secondary mask argument (used for more complex distributions), see help in compute_mask_indices + self.no_mask_channel_overlap: bool = False # whether to allow channel masks to overlap + self.mask_channel_min_space: int = 1 # min space between spans (if no overlap is enabled) + + # positional embeddings + self.conv_pos: int = 128 # number of filters for convolutional positional embeddings + self.conv_pos_groups: int = 16 # number of groups for convolutional positional embedding + + # relative position embedding + self.relative_position_embedding: bool = False # apply relative position embedding + self.num_buckets: int = 320 # number of buckets for relative position embedding + self.max_distance: int = 1280 # maximum distance for relative position embedding + self.gru_rel_pos: bool = False # apply gated relative position embedding + + if cfg is not None: + self.update(cfg) + + def update(self, cfg: dict): + self.__dict__.update(cfg) + + +class WavLM(nn.Module): + def __init__( + self, + cfg: WavLMConfig, + ) -> None: + super().__init__() + logger.info(f"WavLM Config: {cfg.__dict__}") + + self.cfg = cfg + feature_enc_layers = eval(cfg.conv_feature_layers) + self.embed = feature_enc_layers[-1][0] + + self.feature_extractor = ConvFeatureExtractionModel( + conv_layers=feature_enc_layers, + dropout=0.0, + mode=cfg.extractor_mode, + conv_bias=cfg.conv_bias, + ) + + self.post_extract_proj = ( + nn.Linear(self.embed, cfg.encoder_embed_dim) + if self.embed != cfg.encoder_embed_dim + else None + ) + + self.mask_prob = cfg.mask_prob + self.mask_selection = cfg.mask_selection + self.mask_other = cfg.mask_other + self.mask_length = cfg.mask_length + self.no_mask_overlap = cfg.no_mask_overlap + self.mask_min_space = cfg.mask_min_space + + self.mask_channel_prob = cfg.mask_channel_prob + self.mask_channel_selection = cfg.mask_channel_selection + self.mask_channel_other = cfg.mask_channel_other + self.mask_channel_length = cfg.mask_channel_length + self.no_mask_channel_overlap = cfg.no_mask_channel_overlap + self.mask_channel_min_space = cfg.mask_channel_min_space + + self.dropout_input = nn.Dropout(cfg.dropout_input) + self.dropout_features = nn.Dropout(cfg.dropout_features) + + self.feature_grad_mult = cfg.feature_grad_mult + + self.mask_emb = nn.Parameter( + torch.FloatTensor(cfg.encoder_embed_dim).uniform_() + ) + + self.encoder = TransformerEncoder(cfg) + self.layer_norm = LayerNorm(self.embed) + + def apply_mask(self, x, padding_mask): + B, T, C = x.shape + if self.mask_prob > 0: + mask_indices = compute_mask_indices( + (B, T), + padding_mask, + self.mask_prob, + self.mask_length, + self.mask_selection, + self.mask_other, + min_masks=2, + no_overlap=self.no_mask_overlap, + min_space=self.mask_min_space, + ) + mask_indices = torch.from_numpy(mask_indices).to(x.device) + x[mask_indices] = self.mask_emb + else: + mask_indices = None + + if self.mask_channel_prob > 0: + mask_channel_indices = compute_mask_indices( + (B, C), + None, + self.mask_channel_prob, + self.mask_channel_length, + self.mask_channel_selection, + self.mask_channel_other, + no_overlap=self.no_mask_channel_overlap, + min_space=self.mask_channel_min_space, + ) + mask_channel_indices = ( + torch.from_numpy(mask_channel_indices) + .to(x.device) + .unsqueeze(1) + .expand(-1, T, -1) + ) + x[mask_channel_indices] = 0 + + return x, mask_indices + + def forward_padding_mask( + self, features: torch.Tensor, padding_mask: torch.Tensor, + ) -> torch.Tensor: + extra = padding_mask.size(1) % features.size(1) + if extra > 0: + padding_mask = padding_mask[:, :-extra] + padding_mask = padding_mask.view( + padding_mask.size(0), features.size(1), -1 + ) + padding_mask = padding_mask.all(-1) + return padding_mask + + def extract_features( + self, + source: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + mask: bool = False, + ret_conv: bool = False, + output_layer: Optional[int] = None, + ret_layer_results: bool = False, + ): + + + with torch.no_grad(): + features = self.feature_extractor(source) + + cnn_outs = features + features = features[-1].transpose(1, 2) + features = self.layer_norm(features) + + if padding_mask is not None: + padding_mask = self.forward_padding_mask(features, padding_mask) + + if self.post_extract_proj is not None: + features = self.post_extract_proj(features) + + features = self.dropout_input(features) + + if mask: + x, mask_indices = self.apply_mask( + features, padding_mask + ) + else: + x = features + + # feature: (B, T, D), float + # target: (B, T), long + # x: (B, T, D), float + # padding_mask: (B, T), bool + # mask_indices: (B, T), bool + x, layer_results = self.encoder( + x, + padding_mask=padding_mask, + layer=None if output_layer is None else output_layer - 1 + ) + return cnn_outs, layer_results + # res = {"x": x, "padding_mask": padding_mask, "features": features, "layer_results": layer_results} + + # feature = res["features"] if ret_conv else res["x"] + # if ret_layer_results: + # feature = (feature, res["layer_results"]) + # return feature, res["padding_mask"] + + +class ConvFeatureExtractionModel(nn.Module): + def __init__( + self, + conv_layers: List[Tuple[int, int, int]], + dropout: float = 0.0, + mode: str = "default", + conv_bias: bool = False, + conv_type: str = "default" + ): + super().__init__() + + assert mode in {"default", "layer_norm"} + + def block( + n_in, + n_out, + k, + stride, + is_layer_norm=False, + is_group_norm=False, + conv_bias=False, + ): + def make_conv(): + conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) + nn.init.kaiming_normal_(conv.weight) + return conv + + assert ( + is_layer_norm and is_group_norm + ) == False, "layer norm and group norm are exclusive" + + if is_layer_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + nn.Sequential( + TransposeLast(), + Fp32LayerNorm(dim, elementwise_affine=True), + TransposeLast(), + ), + nn.GELU(), + ) + elif is_group_norm: + return nn.Sequential( + make_conv(), + nn.Dropout(p=dropout), + Fp32GroupNorm(dim, dim, affine=True), + nn.GELU(), + ) + else: + return nn.Sequential(make_conv(), nn.Dropout(p=dropout), nn.GELU()) + + self.conv_type = conv_type + if self.conv_type == "default": + in_d = 1 + self.conv_layers = nn.ModuleList() + for i, cl in enumerate(conv_layers): + assert len(cl) == 3, "invalid conv definition: " + str(cl) + (dim, k, stride) = cl + + self.conv_layers.append( + block( + in_d, + dim, + k, + stride, + is_layer_norm=mode == "layer_norm", + is_group_norm=mode == "default" and i == 0, + conv_bias=conv_bias, + ) + ) + in_d = dim + elif self.conv_type == "conv2d": + in_d = 1 + self.conv_layers = nn.ModuleList() + for i, cl in enumerate(conv_layers): + assert len(cl) == 3 + (dim, k, stride) = cl + + self.conv_layers.append( + torch.nn.Conv2d(in_d, dim, k, stride) + ) + self.conv_layers.append(torch.nn.ReLU()) + in_d = dim + elif self.conv_type == "custom": + in_d = 1 + idim = 80 + self.conv_layers = nn.ModuleList() + for i, cl in enumerate(conv_layers): + assert len(cl) == 3 + (dim, k, stride) = cl + self.conv_layers.append( + torch.nn.Conv2d(in_d, dim, k, stride, padding=1) + ) + self.conv_layers.append( + torch.nn.LayerNorm([dim, idim]) + ) + self.conv_layers.append(torch.nn.ReLU()) + in_d = dim + if (i + 1) % 2 == 0: + self.conv_layers.append( + torch.nn.MaxPool2d(2, stride=2, ceil_mode=True) + ) + idim = int(math.ceil(idim / 2)) + else: + pass + + def forward(self, x, mask=None): + + # BxT -> BxCxT + x_lst = [] + x = x.unsqueeze(1) + if self.conv_type == "custom": + for conv in self.conv_layers: + if isinstance(conv, nn.LayerNorm): + x = x.transpose(1, 2) + x = conv(x).transpose(1, 2) + else: + x = conv(x) + x = x.transpose(2, 3).contiguous() + x = x.view(x.size(0), -1, x.size(-1)) + else: + for conv in self.conv_layers: + x = conv(x) + x_lst.append(x) + if self.conv_type == "conv2d": + b, c, t, f = x.size() + x = x.transpose(2, 3).contiguous().view(b, c * f, t) + return x_lst + + +class TransformerEncoder(nn.Module): + def __init__(self, args): + super().__init__() + + self.dropout = args.dropout + self.embedding_dim = args.encoder_embed_dim + + self.pos_conv = nn.Conv1d( + self.embedding_dim, + self.embedding_dim, + kernel_size=args.conv_pos, + padding=args.conv_pos // 2, + groups=args.conv_pos_groups, + ) + dropout = 0 + std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim)) + nn.init.normal_(self.pos_conv.weight, mean=0, std=std) + nn.init.constant_(self.pos_conv.bias, 0) + + self.pos_conv = nn.utils.weight_norm(self.pos_conv, name="weight", dim=2) + self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU()) + + if hasattr(args, "relative_position_embedding"): + self.relative_position_embedding = args.relative_position_embedding + self.num_buckets = args.num_buckets + self.max_distance = args.max_distance + else: + self.relative_position_embedding = False + self.num_buckets = 0 + self.max_distance = 0 + + self.layers = nn.ModuleList( + [ + TransformerSentenceEncoderLayer( + embedding_dim=self.embedding_dim, + ffn_embedding_dim=args.encoder_ffn_embed_dim, + num_attention_heads=args.encoder_attention_heads, + dropout=self.dropout, + attention_dropout=args.attention_dropout, + activation_dropout=args.activation_dropout, + activation_fn=args.activation_fn, + layer_norm_first=args.layer_norm_first, + has_relative_attention_bias=(self.relative_position_embedding and i == 0), + num_buckets=self.num_buckets, + max_distance=self.max_distance, + gru_rel_pos=args.gru_rel_pos, + ) + for i in range(args.encoder_layers) + ] + ) + + self.layer_norm_first = args.layer_norm_first + self.layer_norm = LayerNorm(self.embedding_dim) + self.layerdrop = args.encoder_layerdrop + + self.apply(init_bert_params) + + def forward(self, x, padding_mask=None, streaming_mask=None, layer=None): + x, layer_results = self.extract_features(x, padding_mask, streaming_mask, layer) + + if self.layer_norm_first and layer is None: + x = self.layer_norm(x) + + return x, layer_results + + def extract_features(self, x, padding_mask=None, streaming_mask=None, tgt_layer=None): + + if padding_mask is not None: + x[padding_mask] = 0 + + x_conv = self.pos_conv(x.transpose(1, 2)) + x_conv = x_conv.transpose(1, 2) + x = x + x_conv + + if not self.layer_norm_first: + x = self.layer_norm(x) + + x = F.dropout(x, p=self.dropout, training=self.training) + + # B x T x C -> T x B x C + x = x.transpose(0, 1) + + layer_results = [] + z = None + if tgt_layer is not None: + layer_results.append((x, z)) + r = None + pos_bias = None + for i, layer in enumerate(self.layers): + dropout_probability = np.random.random() + if not self.training or (dropout_probability > self.layerdrop): + x, z, pos_bias = layer(x, self_attn_padding_mask=padding_mask, need_weights=False, + self_attn_mask=streaming_mask, pos_bias=pos_bias) + if tgt_layer is not None: + layer_results.append((x, z)) + if i == tgt_layer: + r = x + break + + if r is not None: + x = r + + # T x B x C -> B x T x C + x = x.transpose(0, 1) + + return x, layer_results + + +class TransformerSentenceEncoderLayer(nn.Module): + """ + Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained + models. + """ + + def __init__( + self, + embedding_dim: float = 768, + ffn_embedding_dim: float = 3072, + num_attention_heads: float = 8, + dropout: float = 0.1, + attention_dropout: float = 0.1, + activation_dropout: float = 0.1, + activation_fn: str = "relu", + layer_norm_first: bool = False, + has_relative_attention_bias: bool = False, + num_buckets: int = 0, + max_distance: int = 0, + rescale_init: bool = False, + gru_rel_pos: bool = False, + ) -> None: + + super().__init__() + # Initialize parameters + self.embedding_dim = embedding_dim + self.dropout = dropout + self.activation_dropout = activation_dropout + + # Initialize blocks + self.activation_name = activation_fn + self.activation_fn = get_activation_fn(activation_fn) + self.self_attn = MultiheadAttention( + self.embedding_dim, + num_attention_heads, + dropout=attention_dropout, + self_attention=True, + has_relative_attention_bias=has_relative_attention_bias, + num_buckets=num_buckets, + max_distance=max_distance, + rescale_init=rescale_init, + gru_rel_pos=gru_rel_pos, + ) + + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(self.activation_dropout) + self.dropout3 = nn.Dropout(dropout) + + self.layer_norm_first = layer_norm_first + + # layer norm associated with the self attention layer + self.self_attn_layer_norm = LayerNorm(self.embedding_dim) + + if self.activation_name == "glu": + self.fc1 = GLU_Linear(self.embedding_dim, ffn_embedding_dim, "swish") + else: + self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim) + self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim) + + import torchaudio.functional as AudioF + + + # layer norm associated with the position wise feed-forward NN + self.final_layer_norm = LayerNorm(self.embedding_dim) + + def forward( + self, + x: torch.Tensor, + self_attn_mask: torch.Tensor = None, + self_attn_padding_mask: torch.Tensor = None, + need_weights: bool = False, + pos_bias=None + ): + """ + LayerNorm is applied either before or after the self-attention/ffn + modules similar to the original Transformer imlementation. + """ + residual = x + + if self.layer_norm_first: + x = self.self_attn_layer_norm(x) + x, attn, pos_bias = self.self_attn( + query=x, + key=x, + value=x, + key_padding_mask=self_attn_padding_mask, + need_weights=False, + attn_mask=self_attn_mask, + position_bias=pos_bias + ) + x = self.dropout1(x) + x = residual + x + + residual = x + + x = self.final_layer_norm(x) + if self.activation_name == "glu": + x = self.fc1(x) + else: + x = self.activation_fn(self.fc1(x)) + x = self.dropout2(x) + x = self.fc2(x) + x = self.dropout3(x) + x = residual + x + else: + + + x, attn, pos_bias = self.self_attn( + query=x, + key=x, + value=x, + key_padding_mask=self_attn_padding_mask, + need_weights=need_weights, + attn_mask=self_attn_mask, + position_bias=pos_bias + ) + + x = self.dropout1(x) + x = residual + x + + x = self.self_attn_layer_norm(x) + + residual = x + if self.activation_name == "glu": + x = self.fc1(x) + else: + x = self.activation_fn(self.fc1(x)) + x = self.dropout2(x) + x = self.fc2(x) + x = self.dropout3(x) + x = residual + x + + + x = self.final_layer_norm(x) + + return x, attn, pos_bias diff --git a/models/Baseline/modules.py b/models/Baseline/modules.py new file mode 100644 index 0000000..1dcfc6f --- /dev/null +++ b/models/Baseline/modules.py @@ -0,0 +1,827 @@ +# -------------------------------------------------------- +# WavLM: Large-Scale Self-Supervised Pre-training for Full Stack Speech Processing (https://arxiv.org/abs/2110.13900.pdf) +# Github source: https://github.com/microsoft/unilm/tree/master/wavlm +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Based on fairseq code bases +# https://github.com/pytorch/fairseq +# -------------------------------------------------------- + +import math +import warnings +from typing import Dict, Optional, Tuple +import torch +from torch import Tensor, nn +from torch.nn import Parameter +import torch.nn.functional as F + + +class TransposeLast(nn.Module): + def __init__(self, deconstruct_idx=None): + super().__init__() + self.deconstruct_idx = deconstruct_idx + + def forward(self, x): + if self.deconstruct_idx is not None: + x = x[self.deconstruct_idx] + return x.transpose(-2, -1) + + +class Fp32LayerNorm(nn.LayerNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.layer_norm( + input.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +class Fp32GroupNorm(nn.GroupNorm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def forward(self, input): + output = F.group_norm( + input.float(), + self.num_groups, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ) + return output.type_as(input) + + +class GradMultiply(torch.autograd.Function): + @staticmethod + def forward(ctx, x, scale): + ctx.scale = scale + res = x.new(x) + return res + + @staticmethod + def backward(ctx, grad): + return grad * ctx.scale, None + + +class SamePad(nn.Module): + def __init__(self, kernel_size, causal=False): + super().__init__() + if causal: + self.remove = kernel_size - 1 + else: + self.remove = 1 if kernel_size % 2 == 0 else 0 + + def forward(self, x): + if self.remove > 0: + x = x[:, :, : -self.remove] + return x + + +class Swish(nn.Module): + """Swish function + """ + + def __init__(self): + """Construct an MultiHeadedAttention object.""" + super(Swish, self).__init__() + self.act = torch.nn.Sigmoid() + + def forward(self, x): + return x * self.act(x) + + +class GLU_Linear(nn.Module): + def __init__(self, input_dim, output_dim, glu_type="sigmoid", bias_in_glu=True): + super(GLU_Linear, self).__init__() + + self.glu_type = glu_type + self.output_dim = output_dim + + if glu_type == "sigmoid": + self.glu_act = torch.nn.Sigmoid() + elif glu_type == "swish": + self.glu_act = Swish() + elif glu_type == "relu": + self.glu_act = torch.nn.ReLU() + elif glu_type == "gelu": + self.glu_act = torch.nn.GELU() + + if bias_in_glu: + self.linear = nn.Linear(input_dim, output_dim * 2, True) + else: + self.linear = nn.Linear(input_dim, output_dim * 2, False) + + def forward(self, x): + # to be consistent with GLU_Linear, we assume the input always has the #channel (#dim) in the last dimension of the tensor, so need to switch the dimension first for 1D-Conv case + x = self.linear(x) + + if self.glu_type == "bilinear": + x = (x[:, :, 0:self.output_dim] * x[:, :, self.output_dim:self.output_dim * 2]) + else: + x = (x[:, :, 0:self.output_dim] * self.glu_act(x[:, :, self.output_dim:self.output_dim * 2])) + + return x + + +def gelu_accurate(x): + if not hasattr(gelu_accurate, "_a"): + gelu_accurate._a = math.sqrt(2 / math.pi) + return ( + 0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3)))) + ) + + +def gelu(x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(x.float()).type_as(x) + + +def get_activation_fn(activation: str): + """Returns the activation function corresponding to `activation`""" + + if activation == "relu": + return F.relu + elif activation == "gelu": + return gelu + elif activation == "gelu_fast": + warnings.warn( + "--activation-fn=gelu_fast has been renamed to gelu_accurate" + ) + return gelu_accurate + elif activation == "gelu_accurate": + return gelu_accurate + elif activation == "tanh": + return torch.tanh + elif activation == "linear": + return lambda x: x + elif activation == "glu": + return lambda x: x + else: + raise RuntimeError("--activation-fn {} not supported".format(activation)) + + +def init_bert_params(module): + """ + Initialize the weights specific to the BERT Model. + This overrides the default initializations depending on the specified arguments. + 1. If normal_init_linear_weights is set then weights of linear + layer will be initialized using the normal distribution and + bais will be set to the specified value. + 2. If normal_init_embed_weights is set then weights of embedding + layer will be initialized using the normal distribution. + 3. If normal_init_proj_weights is set then weights of + in_project_weight for MultiHeadAttention initialized using + the normal distribution (to be validated). + """ + + def normal_(data): + # with FSDP, module params will be on CUDA, so we cast them back to CPU + # so that the RNG is consistent with and without FSDP + data.copy_( + data.cpu().normal_(mean=0.0, std=0.02).to(data.device) + ) + + if isinstance(module, nn.Linear): + normal_(module.weight.data) + if module.bias is not None: + module.bias.data.zero_() + if isinstance(module, nn.Embedding): + normal_(module.weight.data) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + if isinstance(module, MultiheadAttention): + normal_(module.q_proj.weight.data) + normal_(module.k_proj.weight.data) + normal_(module.v_proj.weight.data) + + +def quant_noise(module, p, block_size): + """ + Wraps modules and applies quantization noise to the weights for + subsequent quantization with Iterative Product Quantization as + described in "Training with Quantization Noise for Extreme Model Compression" + + Args: + - module: nn.Module + - p: amount of Quantization Noise + - block_size: size of the blocks for subsequent quantization with iPQ + + Remarks: + - Module weights must have the right sizes wrt the block size + - Only Linear, Embedding and Conv2d modules are supported for the moment + - For more detail on how to quantize by blocks with convolutional weights, + see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks" + - We implement the simplest form of noise here as stated in the paper + which consists in randomly dropping blocks + """ + + # if no quantization noise, don't register hook + if p <= 0: + return module + + # supported modules + assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d)) + + # test whether module.weight has the right sizes wrt block_size + is_conv = module.weight.ndim == 4 + + # 2D matrix + if not is_conv: + assert ( + module.weight.size(1) % block_size == 0 + ), "Input features must be a multiple of block sizes" + + # 4D matrix + else: + # 1x1 convolutions + if module.kernel_size == (1, 1): + assert ( + module.in_channels % block_size == 0 + ), "Input channels must be a multiple of block sizes" + # regular convolutions + else: + k = module.kernel_size[0] * module.kernel_size[1] + assert k % block_size == 0, "Kernel size must be a multiple of block size" + + def _forward_pre_hook(mod, input): + # no noise for evaluation + if mod.training: + if not is_conv: + # gather weight and sizes + weight = mod.weight + in_features = weight.size(1) + out_features = weight.size(0) + + # split weight matrix into blocks and randomly drop selected blocks + mask = torch.zeros( + in_features // block_size * out_features, device=weight.device + ) + mask.bernoulli_(p) + mask = mask.repeat_interleave(block_size, -1).view(-1, in_features) + + else: + # gather weight and sizes + weight = mod.weight + in_channels = mod.in_channels + out_channels = mod.out_channels + + # split weight matrix into blocks and randomly drop selected blocks + if mod.kernel_size == (1, 1): + mask = torch.zeros( + int(in_channels // block_size * out_channels), + device=weight.device, + ) + mask.bernoulli_(p) + mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels) + else: + mask = torch.zeros( + weight.size(0), weight.size(1), device=weight.device + ) + mask.bernoulli_(p) + mask = ( + mask.unsqueeze(2) + .unsqueeze(3) + .repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1]) + ) + + # scale weights and apply mask + mask = mask.to( + torch.bool + ) # x.bool() is not currently supported in TorchScript + s = 1 / (1 - p) + mod.weight.data = s * weight.masked_fill(mask, 0) + + module.register_forward_pre_hook(_forward_pre_hook) + return module + + +class MultiheadAttention(nn.Module): + """Multi-headed attention. + + See "Attention Is All You Need" for more details. + """ + + def __init__( + self, + embed_dim, + num_heads, + kdim=None, + vdim=None, + dropout=0.0, + bias=True, + add_bias_kv=False, + add_zero_attn=False, + self_attention=False, + encoder_decoder_attention=False, + q_noise=0.0, + qn_block_size=8, + has_relative_attention_bias=False, + num_buckets=32, + max_distance=128, + gru_rel_pos=False, + rescale_init=False, + ): + super().__init__() + self.embed_dim = embed_dim + self.kdim = kdim if kdim is not None else embed_dim + self.vdim = vdim if vdim is not None else embed_dim + self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim + + self.num_heads = num_heads + self.dropout_module = nn.Dropout(dropout) + + self.has_relative_attention_bias = has_relative_attention_bias + self.num_buckets = num_buckets + self.max_distance = max_distance + if self.has_relative_attention_bias: + self.relative_attention_bias = nn.Embedding(num_buckets, num_heads) + + self.head_dim = embed_dim // num_heads + self.q_head_dim = self.head_dim + self.k_head_dim = self.head_dim + assert ( + self.head_dim * num_heads == self.embed_dim + ), "embed_dim must be divisible by num_heads" + self.scaling = self.head_dim ** -0.5 + + self.self_attention = self_attention + self.encoder_decoder_attention = encoder_decoder_attention + + assert not self.self_attention or self.qkv_same_dim, ( + "Self-attention requires query, key and " "value to be of the same size" + ) + + k_bias = True + if rescale_init: + k_bias = False + + k_embed_dim = embed_dim + q_embed_dim = embed_dim + + self.k_proj = quant_noise( + nn.Linear(self.kdim, k_embed_dim, bias=k_bias), q_noise, qn_block_size + ) + self.v_proj = quant_noise( + nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size + ) + self.q_proj = quant_noise( + nn.Linear(embed_dim, q_embed_dim, bias=bias), q_noise, qn_block_size + ) + + self.out_proj = quant_noise( + nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size + ) + + if add_bias_kv: + self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim)) + self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim)) + else: + self.bias_k = self.bias_v = None + + self.add_zero_attn = add_zero_attn + + self.gru_rel_pos = gru_rel_pos + if self.gru_rel_pos: + self.grep_linear = nn.Linear(self.q_head_dim, 8) + self.grep_a = nn.Parameter(torch.ones(1, num_heads, 1, 1)) + + self.reset_parameters() + + def reset_parameters(self): + if self.qkv_same_dim: + # Empirically observed the convergence to be much better with + # the scaled initialization + nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2)) + nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2)) + nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2)) + else: + nn.init.xavier_uniform_(self.k_proj.weight) + nn.init.xavier_uniform_(self.v_proj.weight) + nn.init.xavier_uniform_(self.q_proj.weight) + + nn.init.xavier_uniform_(self.out_proj.weight) + if self.out_proj.bias is not None: + nn.init.constant_(self.out_proj.bias, 0.0) + if self.bias_k is not None: + nn.init.xavier_normal_(self.bias_k) + if self.bias_v is not None: + nn.init.xavier_normal_(self.bias_v) + if self.has_relative_attention_bias: + nn.init.xavier_normal_(self.relative_attention_bias.weight) + + def _relative_positions_bucket(self, relative_positions, bidirectional=True): + num_buckets = self.num_buckets + max_distance = self.max_distance + relative_buckets = 0 + + if bidirectional: + num_buckets = num_buckets // 2 + relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets + relative_positions = torch.abs(relative_positions) + else: + relative_positions = -torch.min(relative_positions, torch.zeros_like(relative_positions)) + + max_exact = num_buckets // 2 + is_small = relative_positions < max_exact + + relative_postion_if_large = max_exact + ( + torch.log(relative_positions.float() / max_exact) + / math.log(max_distance / max_exact) + * (num_buckets - max_exact) + ).to(torch.long) + relative_postion_if_large = torch.min( + relative_postion_if_large, torch.full_like(relative_postion_if_large, num_buckets - 1) + ) + + relative_buckets += torch.where(is_small, relative_positions, relative_postion_if_large) + return relative_buckets + + def compute_bias(self, query_length, key_length): + context_position = torch.arange(query_length, dtype=torch.long)[:, None] + memory_position = torch.arange(key_length, dtype=torch.long)[None, :] + relative_position = memory_position - context_position + relative_position_bucket = self._relative_positions_bucket( + relative_position, + bidirectional=True + ) + relative_position_bucket = relative_position_bucket.to(self.relative_attention_bias.weight.device) + values = self.relative_attention_bias(relative_position_bucket) + values = values.permute([2, 0, 1]) + return values + + def forward( + self, + query, + key: Optional[Tensor], + value: Optional[Tensor], + key_padding_mask: Optional[Tensor] = None, + incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, + need_weights: bool = True, + static_kv: bool = False, + attn_mask: Optional[Tensor] = None, + before_softmax: bool = False, + need_head_weights: bool = False, + position_bias: Optional[Tensor] = None + ) -> Tuple[Tensor, Optional[Tensor], Optional[Tensor]]: + """Input shape: Time x Batch x Channel + + Args: + key_padding_mask (ByteTensor, optional): mask to exclude + keys that are pads, of shape `(batch, src_len)`, where + padding elements are indicated by 1s. + need_weights (bool, optional): return the attention weights, + averaged over heads (default: False). + attn_mask (ByteTensor, optional): typically used to + implement causal attention, where the mask prevents the + attention from looking forward in time (default: None). + before_softmax (bool, optional): return the raw attention + weights and values before the attention softmax. + need_head_weights (bool, optional): return the attention + weights for each head. Implies *need_weights*. Default: + return the average attention weights over all heads. + """ + if need_head_weights: + need_weights = True + + is_tpu = query.device.type == "xla" + + tgt_len, bsz, embed_dim = query.size() + src_len = tgt_len + assert embed_dim == self.embed_dim + assert list(query.size()) == [tgt_len, bsz, embed_dim] + if key is not None: + src_len, key_bsz, _ = key.size() + if not torch.jit.is_scripting(): + assert key_bsz == bsz + assert value is not None + assert src_len, bsz == value.shape[:2] + + if self.has_relative_attention_bias and position_bias is None: + position_bias = self.compute_bias(tgt_len, src_len) + position_bias = position_bias.unsqueeze(0).repeat(bsz, 1, 1, 1).view(bsz * self.num_heads, tgt_len, src_len) + + if ( + not is_tpu # don't use PyTorch version on TPUs + and incremental_state is None + and not static_kv + # A workaround for quantization to work. Otherwise JIT compilation + # treats bias in linear module as method. + and not torch.jit.is_scripting() + and self.q_head_dim == self.head_dim + ): + assert key is not None and value is not None + assert attn_mask is None + + attn_mask_rel_pos = None + if position_bias is not None: + attn_mask_rel_pos = position_bias + if self.gru_rel_pos: + query_layer = query.transpose(0, 1) + new_x_shape = query_layer.size()[:-1] + (self.num_heads, -1) + query_layer = query_layer.view(*new_x_shape) + query_layer = query_layer.permute(0, 2, 1, 3) + _B, _H, _L, __ = query_layer.size() + + gate_a, gate_b = torch.sigmoid(self.grep_linear(query_layer).view( + _B, _H, _L, 2, 4).sum(-1, keepdim=False)).chunk(2, dim=-1) + gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0 + attn_mask_rel_pos = gate_a_1.view(bsz * self.num_heads, -1, 1) * position_bias + + attn_mask_rel_pos = attn_mask_rel_pos.view((-1, tgt_len, tgt_len)) + k_proj_bias = self.k_proj.bias + if k_proj_bias is None: + k_proj_bias = torch.zeros_like(self.q_proj.bias) + + x, attn = F.multi_head_attention_forward( + query, + key, + value, + self.embed_dim, + self.num_heads, + torch.empty([0]), + torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)), + self.bias_k, + self.bias_v, + self.add_zero_attn, + self.dropout_module.p, + self.out_proj.weight, + self.out_proj.bias, + self.training, + # self.training or self.dropout_module.apply_during_inference, + key_padding_mask, + need_weights, + attn_mask_rel_pos, + use_separate_proj_weight=True, + q_proj_weight=self.q_proj.weight, + k_proj_weight=self.k_proj.weight, + v_proj_weight=self.v_proj.weight, + ) + return x, attn, position_bias + + if incremental_state is not None: + saved_state = self._get_input_buffer(incremental_state) + if saved_state is not None and "prev_key" in saved_state: + # previous time steps are cached - no need to recompute + # key and value if they are static + if static_kv: + assert self.encoder_decoder_attention and not self.self_attention + key = value = None + else: + saved_state = None + + if self.self_attention: + q = self.q_proj(query) + k = self.k_proj(query) + v = self.v_proj(query) + elif self.encoder_decoder_attention: + # encoder-decoder attention + q = self.q_proj(query) + if key is None: + assert value is None + k = v = None + else: + k = self.k_proj(key) + v = self.v_proj(key) + + else: + assert key is not None and value is not None + q = self.q_proj(query) + k = self.k_proj(key) + v = self.v_proj(value) + q *= self.scaling + + if self.bias_k is not None: + assert self.bias_v is not None + k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)]) + v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)]) + if attn_mask is not None: + attn_mask = torch.cat( + [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1 + ) + if key_padding_mask is not None: + key_padding_mask = torch.cat( + [ + key_padding_mask, + key_padding_mask.new_zeros(key_padding_mask.size(0), 1), + ], + dim=1, + ) + + q = ( + q.contiguous() + .view(tgt_len, bsz * self.num_heads, self.q_head_dim) + .transpose(0, 1) + ) + if k is not None: + k = ( + k.contiguous() + .view(-1, bsz * self.num_heads, self.k_head_dim) + .transpose(0, 1) + ) + if v is not None: + v = ( + v.contiguous() + .view(-1, bsz * self.num_heads, self.head_dim) + .transpose(0, 1) + ) + + if saved_state is not None: + # saved states are stored with shape (bsz, num_heads, seq_len, head_dim) + if "prev_key" in saved_state: + _prev_key = saved_state["prev_key"] + assert _prev_key is not None + prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim) + if static_kv: + k = prev_key + else: + assert k is not None + k = torch.cat([prev_key, k], dim=1) + src_len = k.size(1) + if "prev_value" in saved_state: + _prev_value = saved_state["prev_value"] + assert _prev_value is not None + prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim) + if static_kv: + v = prev_value + else: + assert v is not None + v = torch.cat([prev_value, v], dim=1) + prev_key_padding_mask: Optional[Tensor] = None + if "prev_key_padding_mask" in saved_state: + prev_key_padding_mask = saved_state["prev_key_padding_mask"] + assert k is not None and v is not None + key_padding_mask = MultiheadAttention._append_prev_key_padding_mask( + key_padding_mask=key_padding_mask, + prev_key_padding_mask=prev_key_padding_mask, + batch_size=bsz, + src_len=k.size(1), + static_kv=static_kv, + ) + + saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim) + saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim) + saved_state["prev_key_padding_mask"] = key_padding_mask + # In this branch incremental_state is never None + assert incremental_state is not None + incremental_state = self._set_input_buffer(incremental_state, saved_state) + assert k is not None + assert k.size(1) == src_len + + # This is part of a workaround to get around fork/join parallelism + # not supporting Optional types. + if key_padding_mask is not None and key_padding_mask.dim() == 0: + key_padding_mask = None + + if key_padding_mask is not None: + assert key_padding_mask.size(0) == bsz + assert key_padding_mask.size(1) == src_len + + if self.add_zero_attn: + assert v is not None + src_len += 1 + k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1) + v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1) + if attn_mask is not None: + attn_mask = torch.cat( + [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1 + ) + if key_padding_mask is not None: + key_padding_mask = torch.cat( + [ + key_padding_mask, + torch.zeros(key_padding_mask.size(0), 1).type_as( + key_padding_mask + ), + ], + dim=1, + ) + + attn_weights = torch.bmm(q, k.transpose(1, 2)) + attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz) + + assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len] + + if attn_mask is not None: + attn_mask = attn_mask.unsqueeze(0) + attn_weights += attn_mask + + if key_padding_mask is not None: + # don't attend to padding symbols + attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + if not is_tpu: + attn_weights = attn_weights.masked_fill( + key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), + float("-inf"), + ) + else: + attn_weights = attn_weights.transpose(0, 2) + attn_weights = attn_weights.masked_fill(key_padding_mask, float("-inf")) + attn_weights = attn_weights.transpose(0, 2) + attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) + + if before_softmax: + return attn_weights, v, position_bias + + if position_bias is not None: + if self.gru_rel_pos == 1: + query_layer = q.view(bsz, self.num_heads, tgt_len, self.q_head_dim) + _B, _H, _L, __ = query_layer.size() + gate_a, gate_b = torch.sigmoid(self.grep_linear(query_layer).view( + _B, _H, _L, 2, 4).sum(-1, keepdim=False)).chunk(2, dim=-1) + gate_a_1 = gate_a * (gate_b * self.grep_a - 1.0) + 2.0 + position_bias = gate_a_1.view(bsz * self.num_heads, -1, 1) * position_bias + + position_bias = position_bias.view(attn_weights.size()) + + attn_weights = attn_weights + position_bias + + attn_weights_float = F.softmax( + attn_weights, dim=-1 + ) + attn_weights = attn_weights_float.type_as(attn_weights) + attn_probs = self.dropout_module(attn_weights) + + assert v is not None + attn = torch.bmm(attn_probs, v) + assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim] + attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim) + attn = self.out_proj(attn) + attn_weights: Optional[Tensor] = None + if need_weights: + attn_weights = attn_weights_float.view( + bsz, self.num_heads, tgt_len, src_len + ).transpose(1, 0) + if not need_head_weights: + # average attention weights over heads + attn_weights = attn_weights.mean(dim=0) + + return attn, attn_weights, position_bias + + @staticmethod + def _append_prev_key_padding_mask( + key_padding_mask: Optional[Tensor], + prev_key_padding_mask: Optional[Tensor], + batch_size: int, + src_len: int, + static_kv: bool, + ) -> Optional[Tensor]: + # saved key padding masks have shape (bsz, seq_len) + if prev_key_padding_mask is not None and static_kv: + new_key_padding_mask = prev_key_padding_mask + elif prev_key_padding_mask is not None and key_padding_mask is not None: + new_key_padding_mask = torch.cat( + [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1 + ) + # During incremental decoding, as the padding token enters and + # leaves the frame, there will be a time when prev or current + # is None + elif prev_key_padding_mask is not None: + if src_len > prev_key_padding_mask.size(1): + filler = torch.zeros( + (batch_size, src_len - prev_key_padding_mask.size(1)), + device=prev_key_padding_mask.device, + ) + new_key_padding_mask = torch.cat( + [prev_key_padding_mask.float(), filler.float()], dim=1 + ) + else: + new_key_padding_mask = prev_key_padding_mask.float() + elif key_padding_mask is not None: + if src_len > key_padding_mask.size(1): + filler = torch.zeros( + (batch_size, src_len - key_padding_mask.size(1)), + device=key_padding_mask.device, + ) + new_key_padding_mask = torch.cat( + [filler.float(), key_padding_mask.float()], dim=1 + ) + else: + new_key_padding_mask = key_padding_mask.float() + else: + new_key_padding_mask = prev_key_padding_mask + return new_key_padding_mask + + def _get_input_buffer( + self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] + ) -> Dict[str, Optional[Tensor]]: + result = self.get_incremental_state(incremental_state, "attn_state") + if result is not None: + return result + else: + empty_result: Dict[str, Optional[Tensor]] = {} + return empty_result + + def _set_input_buffer( + self, + incremental_state: Dict[str, Dict[str, Optional[Tensor]]], + buffer: Dict[str, Optional[Tensor]], + ): + return self.set_incremental_state(incremental_state, "attn_state", buffer) + + def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int): + return attn_weights diff --git a/optimizer/adamw.py b/optimizer/adamw.py new file mode 100644 index 0000000..1218cf4 --- /dev/null +++ b/optimizer/adamw.py @@ -0,0 +1,10 @@ +#! /usr/bin/python +# -*- encoding: utf-8 -*- + +import torch + +def Optimizer(parameters, lr, **kwargs): + + print('Initialised Adam optimizer') + + return torch.optim.AdamW(parameters, lr = lr); diff --git a/pseudo_labeling.py b/pseudo_labeling.py new file mode 100644 index 0000000..6f96c97 --- /dev/null +++ b/pseudo_labeling.py @@ -0,0 +1,79 @@ +import time +import argparse + +import torch + +from cuml.cluster import KMeans + +# from sklearn.cluster import KMeans +from sklearn.cluster import AgglomerativeClustering + +from sklearn.metrics import normalized_mutual_info_score + + +def main(args): + # Load embeddings + embeddings_file = torch.load(args.embeddings_file) + files = list(embeddings_file.keys()) + labels = [file.split('/')[-3] for file in files] + embeddings = torch.cat(list(embeddings_file.values())).numpy() + print(f"Embedding shape: {embeddings.shape}") + + # K-Means + print("KMeans...") + kmeans_start_time = time.time() + kmeans = KMeans( + n_clusters=args.n_clusters, + random_state=0, + max_samples_per_batch=1000000, + verbose=True + ).fit(embeddings) + pseudo_labels = kmeans.labels_ + centroids = kmeans.cluster_centers_ + print(f"K-Means duration: {(time.time() - kmeans_start_time)/60:.2f} min") + + # AHC + if args.n_clusters_ahc > 0: + print("AHC...") + ahc_start_time = time.time() + ahc_labels = AgglomerativeClustering( + n_clusters=args.n_clusters_ahc + ).fit_predict(centroids) + pseudo_labels = [ahc_labels[pl] for pl in pseudo_labels] + print(f"AHC duration: {(time.time() - ahc_start_time)/60:.2f} min") + + # Print NMI + nmi_score = normalized_mutual_info_score(labels, pseudo_labels) + print(f"NMI: {nmi_score}") + + # Export pseudo labels + with open(args.output_file, 'w') as f: + for file, pseudo_label in zip(files, pseudo_labels): + f.write(f"{pseudo_label} {file}\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + 'embeddings_file', + help='Path to embeddings file (.pt).' + ) + parser.add_argument( + 'output_file', + help='Path to output file (.txt).' + ) + parser.add_argument( + '--n_clusters', + help='Number of clusters for KMeans.', + type=int, + default=50000 + ) + parser.add_argument( + '--n_clusters_ahc', + help='Number of clusters for Agglomerative Clustering.', + type=int, + default=7500 + ) + args = parser.parse_args() + + main(args) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..2d1a60c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +torch>=1.7.0 +torchaudio>=0.7.0 +numpy +scipy +scikit-learn +pyyaml +soundfile + +--extra-index-url https://pypi.nvidia.com +cuml-cu12==24.8.* \ No newline at end of file diff --git a/scheduler/steplr.py b/scheduler/steplr.py new file mode 100644 index 0000000..12364bc --- /dev/null +++ b/scheduler/steplr.py @@ -0,0 +1,16 @@ +#! /usr/bin/python +# -*- encoding: utf-8 -*- + +import torch + +def Scheduler(optimizer, test_interval, max_epoch, lr_decay, **kwargs): + + sche_fn = torch.optim.lr_scheduler.StepLR(optimizer, step_size=test_interval, gamma=lr_decay) + + lr_step = 'epoch' + + print('Initialised step LR scheduler') + + return sche_fn, lr_step + + diff --git a/tools/rsync_jz.sh b/tools/rsync_jz.sh new file mode 100755 index 0000000..7d1d665 --- /dev/null +++ b/tools/rsync_jz.sh @@ -0,0 +1,23 @@ +source_path="." +target_path="jeanzay:~/wavlm_ssl_sv" + +rsync -azh $source_path $target_path \ + --progress \ + --force \ + --delete \ + --exclude="slurm_*" \ + --exclude="data" \ + --exclude="exp" \ + --keep-dirlinks + +while inotifywait -r -e modify,create,delete $source_path +do + rsync -azh $source_path $target_path \ + --progress \ + --force \ + --delete \ + --exclude="slurm_*" \ + --exclude="data" \ + --exclude="exp" \ + --keep-dirlinks +done diff --git a/trainSpeakerNet.py b/trainSpeakerNet.py new file mode 100644 index 0000000..85dec23 --- /dev/null +++ b/trainSpeakerNet.py @@ -0,0 +1,395 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- + +import sys, time, os, argparse, socket +import yaml +import numpy +import pdb +import torch +import glob +import zipfile +import warnings +import datetime +from tuneThreshold import * +from SpeakerNet import * +from DatasetLoader import * +import torch.distributed as dist +import torch.multiprocessing as mp +from scipy.stats import norm +from sklearn.mixture import GaussianMixture + +## ===== ===== ===== ===== ===== ===== ===== ===== +## Parse arguments +## ===== ===== ===== ===== ===== ===== ===== ===== +# os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3' + +parser = argparse.ArgumentParser(description = "SpeakerNet"); + +parser.add_argument('--config', type=str, default=None, help='Config YAML file'); + +## Data loader +parser.add_argument('--max_frames', type=int, default=200, help='Input length to the network for training'); +parser.add_argument('--eval_frames', type=int, default=300, help='Input length to the network for testing; 0 uses the whole files'); +parser.add_argument('--batch_size', type=int, default=400, help='Batch size, number of speakers per batch'); +parser.add_argument('--max_seg_per_spk', type=int, default=500, help='Maximum number of utterances per speaker per epoch'); +parser.add_argument('--nDataLoaderThread', type=int, default=10, help='Number of loader threads'); +parser.add_argument('--augment', type=bool, default=True, help='Augment input') +parser.add_argument('--seed', type=int, default=20211202, help='Seed for the random number generator'); + + + +## Training details +parser.add_argument('--test_interval', type=int, default=1, help='Test and save every [test_interval] epochs'); +parser.add_argument('--max_epoch', type=int, default=50, help='Maximum number of epochs'); +parser.add_argument('--trainfunc', type=str, default="aamsoftmax", help='Loss function'); + +## Optimizer +parser.add_argument('--optimizer', type=str, default="adamw", help='sgd or adam'); +parser.add_argument('--scheduler', type=str, default="steplr", help='Learning rate scheduler'); +parser.add_argument('--lr', type=float, default=0.001, help='Learning rate'); +parser.add_argument("--lr_decay", type=float, default=0.9, help='Learning rate decay every [test_interval] epochs'); + + +## Pre-trained Transformer Model +parser.add_argument('--pretrained_model_path', type=str, default="None", help='Absolute path to the pre-trained model'); +parser.add_argument('--weight_finetuning_reg', type=float, default=0.001, help='L2 regularization towards the initial pre-trained model'); +parser.add_argument('--LLRD_factor', type=float, default=1.0, help='Layer-wise Learning Rate Decay (LLRD) factor'); +parser.add_argument('--LR_Transformer', type=float, default=2e-5, help='Learning rate of pre-trained model'); +parser.add_argument('--LR_MHFA', type=float, default=5e-3, help='Learning rate of back-end attentive pooling model'); + +## Loss functions +parser.add_argument("--hard_prob", type=float, default=0.5, help='Hard negative mining probability, otherwise random, only for some loss functions'); +parser.add_argument("--hard_rank", type=int, default=10, help='Hard negative mining rank in the batch, only for some loss functions'); +parser.add_argument('--margin', type=float, default=0.2, help='Loss margin, only for some loss functions'); +parser.add_argument('--scale', type=float, default=30, help='Loss scale, only for some loss functions'); +parser.add_argument('--nPerSpeaker', type=int, default=1, help='Number of utterances per speaker per batch, only for metric learning based losses'); +parser.add_argument('--nClasses', type=int, default=5994, help='Number of speakers in the softmax layer, only for softmax-based losses'); + +## Evaluation parameters +parser.add_argument('--dcf_p_target', type=float, default=0.05, help='A priori probability of the specified target speaker'); +parser.add_argument('--dcf_c_miss', type=float, default=1, help='Cost of a missed detection'); +parser.add_argument('--dcf_c_fa', type=float, default=1, help='Cost of a spurious detection'); + +## Load and save +parser.add_argument('--initial_model', type=str, default="", help='Initial model weights'); +parser.add_argument('--save_path', type=str, default="exps/exp1", help='Path for model and logs'); + +## Training and test data +parser.add_argument('--train_list', type=str, default="data/train_list.txt", help='Train list'); +parser.add_argument('--test_list', type=str, default="data/test_list.txt", help='Evaluation list'); +parser.add_argument('--train_path', type=str, default="data/voxceleb2", help='Absolute path to the train set'); +parser.add_argument('--test_path', type=str, default="data/voxceleb1", help='Absolute path to the test set'); +parser.add_argument('--musan_path', type=str, default="data/musan_split", help='Absolute path to the test set'); +parser.add_argument('--rir_path', type=str, default="data/simulated_rirs", help='Absolute path to the test set'); + +## Model definition +parser.add_argument('--n_mels', type=int, default=80, help='Number of mel filterbanks'); +parser.add_argument('--log_input', type=bool, default=False, help='Log input features') +parser.add_argument('--model', type=str, default="", help='Name of model definition'); +parser.add_argument('--encoder_type', type=str, default="SAP", help='Type of encoder'); +parser.add_argument('--nOut', type=int, default=192, help='Embedding size in the last FC layer'); + +## For test only +parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only') + +## Distributed and mixed precision training +parser.add_argument('--port', type=str, default="7888", help='Port for distributed training, input as text'); +parser.add_argument('--distributed', dest='distributed', action='store_true', help='Enable distributed training') +parser.add_argument('--mixedprec', dest='mixedprec', action='store_true', help='Enable mixed precision training') + +args = parser.parse_args(); + +## Parse YAML +def find_option_type(key, parser): + for opt in parser._get_optional_actions(): + if ('--' + key) in opt.option_strings: + return opt.type + raise ValueError + +if args.config is not None: + with open(args.config, "r") as f: + yml_config = yaml.load(f, Loader=yaml.FullLoader) + for k, v in yml_config.items(): + if k in args.__dict__: + typ = find_option_type(k, parser) + args.__dict__[k] = typ(v) + else: + sys.stderr.write("Ignored unknown parameter {} in yaml.\n".format(k)) + + +## Try to import NSML +try: + import nsml + from nsml import HAS_DATASET, DATASET_PATH, PARALLEL_WORLD, PARALLEL_PORTS, MY_RANK + from nsml import NSML_NFS_OUTPUT, SESSION_NAME +except: + pass; + +warnings.simplefilter("ignore") + +## ===== ===== ===== ===== ===== ===== ===== ===== +## Trainer script +## ===== ===== ===== ===== ===== ===== ===== ===== + +def LGL_threshold_update_gmm(loss_vals_path): + with open(loss_vals_path, 'r') as f: + lines = [line.strip().split() for line in f.readlines()] + + # losses = [float(line[0]) for line in lines] + losses = [] + errs = 0 + for line in lines: + try: + losses.append(float(line[0])) + except ValueError: + errs += 1 + pass + if errs > 0: + print('Could not read %d lines' % errs) + + log_losses = np.log(losses) + + gmm = GaussianMixture(n_components=2, random_state=0, covariance_type='full', tol=0.00001, max_iter=1000) + gmm.fit(log_losses.reshape(-1, 1)) + + mean1 = gmm.means_[0, 0] + covar1 = gmm.covariances_[0, 0] + weight1 = gmm.weights_[0] + x = np.linspace(min(log_losses), max(log_losses), 1000) + g1 = weight1 * norm.pdf(x, mean1, np.sqrt(covar1)) + + mean2 = gmm.means_[1, 0] + covar2 = gmm.covariances_[1, 0] + weight2 = gmm.weights_[1] + g2 = weight2 * norm.pdf(x, mean2, np.sqrt(covar2)) + + intersection = np.argwhere(np.diff(np.sign(g1 - g2))).flatten() + + max1 = x[np.argmax(g1)] + max2 = x[np.argmax(g2)] + good_intersection = x[intersection][(x[intersection] > min(max1, max2)) & (x[intersection] < max(max1, max2))] + assert len(good_intersection) == 1, 'Wrong number of intersections' + good_intersection = good_intersection[0] + + return good_intersection + +import idr_torch + +def main_worker(gpu, ngpus_per_node, args): + + args.gpu = gpu + + args.gpu = idr_torch.rank + ngpus_per_node = idr_torch.size + + ## Load models + s = SpeakerNet(**vars(args)); + + if args.distributed: + # os.environ['MASTER_ADDR']='localhost' + # os.environ['MASTER_PORT']=args.port + + # dist.init_process_group(backend='nccl', world_size=ngpus_per_node, rank=args.gpu, init_method='tcp://localhost:12345') + dist.init_process_group(backend='nccl', world_size=ngpus_per_node, rank=args.gpu) + + torch.cuda.set_device(args.gpu) + s.cuda(args.gpu) + + s = torch.nn.parallel.DistributedDataParallel(s, device_ids=[args.gpu])#, find_unused_parameters=True) + + print('Loaded the model on GPU {:d}'.format(args.gpu)) + + else: + s = WrappedModel(s).cuda(args.gpu) + + it = 1 + eers = [100]; + + if args.gpu == 0: + ## Write args to scorefile + scorefile = open(args.result_save_path+"/scores.txt", "a+"); + + ## Initialise trainer and data loader + train_dataset = train_dataset_loader(**vars(args)) + + train_sampler = train_dataset_sampler(train_dataset, **vars(args)) + + train_loader = torch.utils.data.DataLoader( + train_dataset, + batch_size=args.batch_size, + num_workers=args.nDataLoaderThread, + sampler=train_sampler, + pin_memory=True, + worker_init_fn=worker_init_fn, + drop_last=True, + ) + + # trainLoader = get_data_loader(args.train_list, **vars(args)); + trainer = ModelTrainer(s, **vars(args)) + + ## Load model weights + modelfiles = glob.glob('%s/model0*.model'%args.model_save_path) + modelfiles.sort() + + if(args.initial_model != ""): + trainer.loadParameters(args.initial_model); + print("Model {} loaded!".format(args.initial_model)); + elif len(modelfiles) >= 1: + print("Model {} loaded from previous state!".format(modelfiles[-1])); + trainer.loadParameters(modelfiles[-1]); + it = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1 + + for ii in range(1,it): + trainer.__scheduler__.step() + + + pytorch_total_params = sum(p.numel() for p in s.module.__S__.parameters()) + + print('Total parameters: ',pytorch_total_params) + ## Evaluation code - must run on single GPU + if args.eval == True: + + + print('Test list',args.test_list) + + sc, lab, _, sc1,sc2 = trainer.evaluateFromList(**vars(args)) + + if args.gpu == 0: + + result = tuneThresholdfromScore(sc, lab, [1, 0.1]); + result_s1 = tuneThresholdfromScore(sc1, lab, [1, 0.1]); + result_s2 = tuneThresholdfromScore(sc2, lab, [1, 0.1]); + + + + fnrs, fprs, thresholds = ComputeErrorRates(sc, lab) + mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, args.dcf_p_target, args.dcf_c_miss, args.dcf_c_fa) + + print('\n',time.strftime("%Y-%m-%d %H:%M:%S"), "VEER {:2.4f}".format(result[1]), "VEER_s1 {:2.4f}".format(result_s1[1]),"VEER_s2 {:2.4f}".format(result_s2[1]),"MinDCF {:2.5f}".format(mindcf)); + + if ("nsml" in sys.modules) and args.gpu == 0: + training_report = {}; + training_report["summary"] = True; + training_report["epoch"] = it; + training_report["step"] = it; + training_report["val_eer"] = result[1]; + training_report["val_dcf"] = mindcf; + + nsml.report(**training_report); + + return + + ## Save training code and params + if args.gpu == 0: + pyfiles = glob.glob('./*.py') + strtime = datetime.datetime.now().strftime("%Y%m%d%H%M%S") + + zipf = zipfile.ZipFile(args.result_save_path+ '/run%s.zip'%strtime, 'w', zipfile.ZIP_DEFLATED) + for file in pyfiles: + zipf.write(file) + zipf.close() + + with open(args.result_save_path + '/run%s.cmd'%strtime, 'w') as f: + f.write('%s'%args) + + + ## Core training script + for it in range(it,args.max_epoch+1): + + train_sampler.set_epoch(it) + + clr = [x['lr'] for x in trainer.__optimizer__.param_groups] + + loss_vals_dir = 'exp/' + args.save_path.split('/')[-1] + '/loss_vals' + os.makedirs(loss_vals_dir, exist_ok=True) + loss_vals_path = os.path.join(loss_vals_dir, 'epoch_%d.txt' % it) + + if it >= 5: + prev_loss_vals_path = os.path.join(loss_vals_dir, 'epoch_%d.txt' % (it - 1)) + LGL_threshold = LGL_threshold_update_gmm(prev_loss_vals_path) + # LGL_threshold = 1 + + if args.gpu == 0: + if LGL_threshold is not None: + print('Updated LGL threshold to %f' % LGL_threshold) + else: + print('Wrong number of intersections, keeping LGL threshold at %f' % LGL_threshold) + + trainer.update_lgl_threshold(LGL_threshold) + + + loss, traineer = trainer.train_network(train_loader, loss_vals_path, it, verbose=(args.gpu == 0)) + + if args.distributed: + dist.barrier() + with open(loss_vals_path, 'w') as final_file: + for r in range(dist.get_world_size()): + part_file_path = f"{loss_vals_path.split('.')[0]}_rank{r}.txt" + with open(part_file_path, 'r') as part_file: + final_file.write(part_file.read()) + + if args.gpu == 0: + print('\n',time.strftime("%Y-%m-%d %H:%M:%S"), "Epoch {:d}, TEER/TAcc {:2.2f}, TLOSS {:f}, LR {:f}".format(it, traineer.item(), loss.item(), max(clr))); + scorefile.write("Epoch {:d}, TEER/TAcc {:2.2f}, TLOSS {:f}, LR {:f} \n".format(it, traineer.item(), loss.item(), max(clr))); + + if it % args.test_interval == 0: + + # sc, lab, _, as1, as2 = trainer.evaluateFromList(**vars(args)) + + if args.gpu == 0: + trainer.saveParameters(args.model_save_path+"/model%09d.model"%it); + + scorefile.flush() + + if ("nsml" in sys.modules) and args.gpu == 0: + training_report = {}; + training_report["summary"] = True; + training_report["epoch"] = it; + training_report["step"] = it; + training_report["train_loss"] = loss; + training_report["min_eer"] = min(eers); + + nsml.report(**training_report); + + if args.gpu == 0: + scorefile.close(); + +## ===== ===== ===== ===== ===== ===== ===== ===== +## Main function +## ===== ===== ===== ===== ===== ===== ===== ===== + + +def main(): + + # print(os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')) + # os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' + # print(os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')) + + + if ("nsml" in sys.modules) and not args.eval: + args.save_path = os.path.join(args.save_path,SESSION_NAME.replace('/','_')) + + args.model_save_path = args.save_path+"/model" + args.result_save_path = args.save_path+"/result" + args.feat_save_path = "" + + os.makedirs(args.model_save_path, exist_ok=True) + os.makedirs(args.result_save_path, exist_ok=True) + + n_gpus = torch.cuda.device_count() + print(n_gpus) + + print('Python Version:', sys.version) + print('PyTorch Version:', torch.__version__) + print('Number of GPUs:', torch.cuda.device_count()) + print('Save path:',args.save_path) + + if args.distributed: + # mp.spawn(main_worker, nprocs=n_gpus, args=(n_gpus, args)) + main_worker(None, None, args) + else: + main_worker(0, None, args) + + +if __name__ == '__main__': + main() diff --git a/trainSpeakerNet_Eval.py b/trainSpeakerNet_Eval.py new file mode 100644 index 0000000..fb5d196 --- /dev/null +++ b/trainSpeakerNet_Eval.py @@ -0,0 +1,250 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- + +import sys, time, os, argparse, socket +import yaml +import numpy +import pdb +import torch +import glob +import zipfile +import csv +import warnings +import datetime +from tuneThreshold import * +from SpeakerNet import * +from DatasetLoader import * +import torch.distributed as dist +import torch.multiprocessing as mp + +## ===== ===== ===== ===== ===== ===== ===== ===== +## Parse arguments +## ===== ===== ===== ===== ===== ===== ===== ===== +# os.environ['CUDA_VISIBLE_DEVICES']='0' +parser = argparse.ArgumentParser(description = "SpeakerNet"); + +parser.add_argument('--config', type=str, default=None, help='Config YAML file'); + +## Data loader +parser.add_argument('--max_frames', type=int, default=200, help='Input length to the network for training'); +parser.add_argument('--eval_frames', type=int, default=300, help='Input length to the network for testing; 0 uses the whole files'); +parser.add_argument('--batch_size', type=int, default=400, help='Batch size, number of speakers per batch'); +parser.add_argument('--max_seg_per_spk', type=int, default=500, help='Maximum number of utterances per speaker per epoch'); +parser.add_argument('--nDataLoaderThread', type=int, default=10, help='Number of loader threads'); +parser.add_argument('--augment', type=bool, default=True, help='Augment input') +parser.add_argument('--seed', type=int, default=20211202, help='Seed for the random number generator'); + + + +## Training details +parser.add_argument('--test_interval', type=int, default=1, help='Test and save every [test_interval] epochs'); +parser.add_argument('--max_epoch', type=int, default=50, help='Maximum number of epochs'); +parser.add_argument('--trainfunc', type=str, default="aamsoftmax", help='Loss function'); + +## Optimizer +parser.add_argument('--optimizer', type=str, default="adamw", help='sgd or adam'); +parser.add_argument('--scheduler', type=str, default="steplr", help='Learning rate scheduler'); +parser.add_argument('--lr', type=float, default=0.001, help='Learning rate'); + + +## Pre-trained Transformer Model +parser.add_argument('--pretrained_model_path', type=str, default="None", help='Absolute path to the pre-trained model'); +parser.add_argument('--weight_finetuning_reg', type=float, default=0.001, help='L2 regularization towards the initial pre-trained model'); +parser.add_argument('--LLRD_factor', type=float, default=1.0, help='Layer-wise Learning Rate Decay (LLRD) factor'); +parser.add_argument('--LR_Transformer', type=float, default=2e-5, help='Learning rate of pre-trained model'); +parser.add_argument('--LR_MHFA', type=float, default=5e-3, help='Learning rate of back-end attentive pooling model'); + +## Loss functions +parser.add_argument("--hard_prob", type=float, default=0.5, help='Hard negative mining probability, otherwise random, only for some loss functions'); +parser.add_argument("--hard_rank", type=int, default=10, help='Hard negative mining rank in the batch, only for some loss functions'); +parser.add_argument('--margin', type=float, default=0.2, help='Loss margin, only for some loss functions'); +parser.add_argument('--scale', type=float, default=30, help='Loss scale, only for some loss functions'); +parser.add_argument('--nPerSpeaker', type=int, default=1, help='Number of utterances per speaker per batch, only for metric learning based losses'); +parser.add_argument('--nClasses', type=int, default=5994, help='Number of speakers in the softmax layer, only for softmax-based losses'); + +## Evaluation parameters +parser.add_argument('--dcf_p_target', type=float, default=0.05, help='A priori probability of the specified target speaker'); +parser.add_argument('--dcf_c_miss', type=float, default=1, help='Cost of a missed detection'); +parser.add_argument('--dcf_c_fa', type=float, default=1, help='Cost of a spurious detection'); + +## Load and save +parser.add_argument('--initial_model', type=str, default="", help='Initial model weights'); +parser.add_argument('--save_path', type=str, default="exps/exp1", help='Path for model and logs'); + +## Training and test data +parser.add_argument('--train_list', type=str, default="data/train_list.txt", help='Train list'); +parser.add_argument('--test_list', type=str, default="data/test_list.txt", help='Evaluation list'); +parser.add_argument('--train_path', type=str, default="data/voxceleb2", help='Absolute path to the train set'); +parser.add_argument('--test_path', type=str, default="data/voxceleb1", help='Absolute path to the test set'); +parser.add_argument('--musan_path', type=str, default="data/musan_split", help='Absolute path to the test set'); +parser.add_argument('--rir_path', type=str, default="data/simulated_rirs", help='Absolute path to the test set'); + +## Model definition +parser.add_argument('--n_mels', type=int, default=80, help='Number of mel filterbanks'); +parser.add_argument('--log_input', type=bool, default=False, help='Log input features') +parser.add_argument('--model', type=str, default="", help='Name of model definition'); +parser.add_argument('--encoder_type', type=str, default="SAP", help='Type of encoder'); +parser.add_argument('--nOut', type=int, default=192, help='Embedding size in the last FC layer'); + +## For test only +parser.add_argument('--eval', dest='eval', action='store_true', help='Eval only') + +parser.add_argument('--generate_embeddings', dest='generate_embeddings', action='store_true', help='Generate embeddings for the train set') +parser.add_argument('--embeddings_path', type=str, default="") +parser.add_argument('--generate_pseudo_labels', dest='generate_pseudo_labels', action='store_true', help='Generate pseudo labels for the train set') + +## Distributed and mixed precision training +parser.add_argument('--port', type=str, default="7888", help='Port for distributed training, input as text'); +parser.add_argument('--distributed', dest='distributed', action='store_true', help='Enable distributed training') +parser.add_argument('--mixedprec', dest='mixedprec', action='store_true', help='Enable mixed precision training') + +args = parser.parse_args(); + +## Parse YAML +def find_option_type(key, parser): + for opt in parser._get_optional_actions(): + if ('--' + key) in opt.option_strings: + return opt.type + raise ValueError + +if args.config is not None: + with open(args.config, "r") as f: + yml_config = yaml.load(f, Loader=yaml.FullLoader) + for k, v in yml_config.items(): + if k in args.__dict__: + typ = find_option_type(k, parser) + args.__dict__[k] = typ(v) + else: + sys.stderr.write("Ignored unknown parameter {} in yaml.\n".format(k)) + + +## Try to import NSML +try: + import nsml + from nsml import HAS_DATASET, DATASET_PATH, PARALLEL_WORLD, PARALLEL_PORTS, MY_RANK + from nsml import NSML_NFS_OUTPUT, SESSION_NAME +except: + pass; + +warnings.simplefilter("ignore") + +## ===== ===== ===== ===== ===== ===== ===== ===== +## Trainer script +## ===== ===== ===== ===== ===== ===== ===== ===== + +def main_worker(gpu, ngpus_per_node, args): + + args.gpu = gpu + + ## Load models + s = SpeakerNet(**vars(args)); + + + s = WrappedModel(s).cuda(args.gpu) + + it = 1 + eers = [100]; + + # trainLoader = get_data_loader(args.train_list, **vars(args)); + trainer = ModelTrainer(s, **vars(args)) + + ## Load model weights + modelfiles = glob.glob('%s/model0*.model'%args.model_save_path) + modelfiles.sort() + + if(args.initial_model != ""): + trainer.loadParameters(args.initial_model); + print("Model {} loaded!".format(args.initial_model)); + elif len(modelfiles) >= 1: + # print("Model {} loaded from previous state!".format(modelfiles[-2])); + # trainer.loadParameters(modelfiles[-2]); + it = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][5:]) + 1 + + for ii in range(1,it): + trainer.__scheduler__.step() + + + # pytorch_total_params = sum(p.numel() for p in s.module.__S__.model.feature_extractor.parameters()) + pytorch_total_params = sum(p.numel() for p in s.module.__S__.parameters()) + + + print('Total parameters: ',pytorch_total_params) + # quit(); + ## Evaluation code - must run on single GPU + if args.eval == True: + scorefile_score = open(args.result_save_path+"/Eval_scores_mean_O_All.txt", "w"); + print('Test list',args.test_list) + + for i in range(1,15): + print("Model {} loaded from previous state!".format(modelfiles[-i])); + trainer.loadParameters(modelfiles[-i]); + # trainer.loadParameters(modelfiles[0]); + + # sc, lab, _,sc1,sc2 = trainer.evaluateFromList_1utterance(**vars(args)) + sc, lab, _,sc1,sc2 = trainer.evaluateFromList(**vars(args)) + + if args.gpu == 0: + + result = tuneThresholdfromScore(sc, lab, [1, 0.1]); + result1 = tuneThresholdfromScore(sc1, lab, [1, 0.1]); + result2 = tuneThresholdfromScore(sc2, lab, [1, 0.1]); + + fnrs, fprs, thresholds = ComputeErrorRates(sc, lab) + + mindcf, threshold = ComputeMinDcf(fnrs, fprs, thresholds, args.dcf_p_target, args.dcf_c_miss, args.dcf_c_fa) + mindcf_1, threshold_1 = ComputeMinDcf(fnrs, fprs, thresholds, 0.01, args.dcf_c_miss, args.dcf_c_fa) + + print('\n',time.strftime("%Y-%m-%d %H:%M:%S"), "VEER {:2.4f}".format(result[1]),"MinDCF05 {:2.5f}".format(mindcf), "MinDCF01 {:2.5f}".format(mindcf_1)); + + scorefile_score.write("Epoch {}, VEER {:2.4f}, VEER_S1 {:2.4f}, VEER_S2 {:2.4f}, MinDCF05 {:2.5f}, MinDCF01 {:2.5f}\n".format(modelfiles[-i], result[1], result1[1], result2[1], mindcf,mindcf_1)); + scorefile_score.flush() + + scorefile_score.close() + return + + if args.generate_embeddings == True: + print('Generate embeddings for the train set') + wav_list_file = args.train_list + with open(wav_list_file,'r') as f: + wav_files = [args.train_path + '/' + line.strip().split()[1] for line in f.readlines()] + + print("Model {} loaded from previous state!".format(modelfiles[-1])); + trainer.loadParameters(modelfiles[-1]); + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + trainer.generate_embeddings(wav_files, args.embeddings_path, device) + + +## ===== ===== ===== ===== ===== ===== ===== ===== +## Main function +## ===== ===== ===== ===== ===== ===== ===== ===== + + +def main(): + + if ("nsml" in sys.modules) and not args.eval: + args.save_path = os.path.join(args.save_path,SESSION_NAME.replace('/','_')) + + args.model_save_path = args.save_path+"/model" + args.result_save_path = args.save_path+"/result" + args.feat_save_path = "" + + os.makedirs(args.model_save_path, exist_ok=True) + os.makedirs(args.result_save_path, exist_ok=True) + + n_gpus = torch.cuda.device_count() + + print('Python Version:', sys.version) + print('PyTorch Version:', torch.__version__) + print('Number of GPUs:', torch.cuda.device_count()) + print('Save path:',args.save_path) + + if args.distributed: + mp.spawn(main_worker, nprocs=n_gpus, args=(n_gpus, args)) + else: + main_worker(0, None, args) + + +if __name__ == '__main__': + main() diff --git a/train_ddp_jz.sh b/train_ddp_jz.sh new file mode 100644 index 0000000..79213a9 --- /dev/null +++ b/train_ddp_jz.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +#SBATCH --job-name=wavlm_ssl_sv +#SBATCH --output=slurm_%j +#SBATCH --nodes=1 +#SBATCH --ntasks=2 +#SBATCH --gres=gpu:2 +#SBATCH --cpus-per-task=10 +#SBATCH --constraint=a100 +#SBATCH --time=20:00:00 +#SBATCH --hint=nomultithread +#SBATCH --account=kdp@a100 + +module purge + +module load cpuarch/amd +module load pytorch-gpu/py3/1.12.1 + +srun python -u trainSpeakerNet.py --config configs/wavlm_mhfa_dlg_lc.yaml --train_list exp/train_list_dino.txt --distributed \ No newline at end of file diff --git a/training_framework.svg b/training_framework.svg new file mode 100755 index 0000000..6eae1f7 --- /dev/null +++ b/training_framework.svg @@ -0,0 +1,3 @@ + + +
Self-Supervised
DINO Training
Pseudo-Labels
Creation
WavLM MHFA
Fine-Tuning
Large Margin
Fine-Tuning
Iterative (×2)
k-means
(50,000)
AHC
(7,500)
Dynamic Loss Gate
Linear
AAM Softmax
Label Correction
Linear
Linear
Concat
WavLM (pre-trained)
Key flow
Value flow
CNN Encoder
Pseudo-labels
Transformer Layer L
Transformer Layer 1
. . .
MHFA
Original
+ aug.
frames
Attentive Pooling
Speaker Embeddings
Reliable
labels
Unreliable
labels
Step 1
Step 2
Step 3
Step 4
Encoder
Projector
Student branch
Encoder
Projector
EMA
DINO loss
Teacher branch
DINO
Redundancy elimination and diversity regularization
4 short aug. frames
2 long aug. frames
Speaker Embeddings
Speaker Embeddings
 \ No newline at end of file diff --git a/tuneThreshold.py b/tuneThreshold.py new file mode 100644 index 0000000..aee3cdc --- /dev/null +++ b/tuneThreshold.py @@ -0,0 +1,87 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- + +import os +import glob +import sys +import time +from sklearn import metrics +import numpy +import pdb +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] + 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): + + # Sort the scores from smallest to largest, and also get the corresponding + # indexes of the sorted scores. We will treat the sorted scores as the + # thresholds at which the the error-rates are evaluated. + sorted_indexes, thresholds = zip(*sorted( + [(index, threshold) for index, threshold in enumerate(scores)], + key=itemgetter(1))) + sorted_labels = [] + labels = [labels[i] for i in sorted_indexes] + fnrs = [] + fprs = [] + + # At the end of this loop, fnrs[i] is the number of errors made by + # incorrectly rejecting scores less than thresholds[i]. And, fprs[i] + # is the total number of times that we have correctly accepted scores + # greater than thresholds[i]. + 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 + + # Now divide by the total number of false negative errors to + # obtain the false positive rates across all thresholds + fnrs = [x / float(fnrs_norm) for x in fnrs] + + # Divide by the total number of corret positives to get the + # true positive rate. Subtract these quantities from 1 to + # get the false positive rates. + 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)): + # See Equation (2). it is a weighted sum of false negative + # and false positive errors. + 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] + # See Equations (3) and (4). Now we normalize the cost. + 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 \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..3d236fa --- /dev/null +++ b/utils.py @@ -0,0 +1,38 @@ +#! /usr/bin/python +# -*- encoding: utf-8 -*- + +import torch +import torch.nn.functional as F + +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + +class PreEmphasis(torch.nn.Module): + + def __init__(self, coef: float = 0.97): + super().__init__() + self.coef = coef + # make kernel + # In pytorch, the convolution operation uses cross-correlation. So, filter is flipped. + self.register_buffer( + 'flipped_filter', torch.FloatTensor([-self.coef, 1.]).unsqueeze(0).unsqueeze(0) + ) + + def forward(self, input: torch.tensor) -> torch.tensor: + assert len(input.size()) == 2, 'The number of dimensions of input tensor must be 2!' + # reflect padding to match lengths of in/out + input = input.unsqueeze(1) + input = F.pad(input, (1, 0), 'reflect') + return F.conv1d(input, self.flipped_filter).squeeze(1) \ No newline at end of file