initial commit

.gitignore

@@ -0,0 +1,6 @@
+__pycache__
+
+exp/
+data/
+
+WavLM-Base+.pt

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()

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.
+
+<p align="center">
+  <img src="training_framework.svg" width=900 />
+</p>
+
+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).

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);
+            
+
+
+
+

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

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

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
+"""

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

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

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

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);

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)

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.*

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
+
+

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

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()

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()

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

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

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)