Pytorch implementation of Diff-SV: A Unified Hierarchical Framework for Noise-Robust Speaker Verification Using Score-Based Diffusion Probabilistic Models
import os
from itertools import chain
import torch

def get_args():
    system_args = {
        # expeirment info
        'project'       : 'NoiseRobustSV',
        'name'          : 'diff-sv',
        'tags'          : ['DFSV'],
        'description'   : 'train and test diff-sv',

        # local
        'path_logging'  : '/results',

        # wandb
        'wandb_group'         : None, 
        'wandb_entity'        : None, 

        # VoxCeleb1 DB
        'path_vox1_train'   : '/data/VoxCeleb1/train',
        'path_vox1_test'    : '/data/VoxCeleb1/test',
        'path_vox1_trials'  : '/data/VoxCeleb1/trials.txt',

        # noise DB
        'path_musan'        : '/data/musan',
        'path_nonspeech'    : '/data/Nonspeech',
        'path_rir'          : '/data/rir_noises/simulated_rirs/',

        # VOiCES DB
        'path_voices_dev'   : '/data/VOiCES_Box_unzip/Development_Data/Speaker_Recognition/sid_dev',
        'path_voices_dev_trials': '/data/VOiCES_Box_unzip/Development_Data/Speaker_Recognition/sid_dev_lists_and_keys/dev-trial-keys.lst',
        'path_voices_eval': '/data/VOiCES_Box_unzip/Speaker_Recognition/sid_eval/',
        'path_voices_eval_trials': '/data/VOiCES_Box_unzip/VOiCES_challenge_2019_post-eval-release/VOiCES_challenge_2019_eval.SID.trial-keys.lst',

        # device
        'num_workers'   : 12,
        'usable_gpu'    : '0,1,2,3',
        'tqdm_ncols'    : 90,
        'path_scripts'     : os.path.dirname(os.path.realpath(__file__))
    }
    
    experiment_args = {
        # env
        'epoch'                     : 320,
        'batch_size'                : 160,
        'number_iteration_for_log'  : 20,
        'rand_seed'                 : 777,
        'flag_reproduciable'        : False,

        # optimizer
        'amsgrad'                   : True,
        'lr_start'                  : 1e-7,
        'lr_end'                    : 1e-7,
        'number_cycle'              : 40,
        'warmup_number_cycle'       : 1,
        'T_mult'                    : 1.5,
        'eta_max'                   : 1e-2,
        'gamma'                     : 0.5,
        'weigth_decay'              : 1e-4,

        # criterion
        'classification_loss'                               : 'aam_softmax',

        # model
        'first_kernel_size'     : 7,
        'first_stride_size'     : (2,1),
        'first_padding_size'    : 3,
        'l_channel'             : [16, 32, 64, 128],
        'l_num_convblocks'      : [3, 4, 6, 3],
        'code_dim'              : 128,
        'stride'                : [1,2,2,1],

        # data
        'winlen'            : 400,
        'winstep'           : 160,
        'train_frame'       : 198,
        'nfft'              : 512,
        'samplerate'        : 16000,
        'nfilts'            : 80,
        'premphasis'        : 0.97,
        'f_min'             : 20,
        'f_max'             : 7600,
        'winfunc'           : torch.hamming_window,
        'DA_p'              : 1,
        'test_frame'        : 998,
        'num_seg'           : 4
    }

    # set args (system_args + experiment_args)
    args = {}
    for k, v in chain(system_args.items(), experiment_args.items()):
        args[k] = v

    return args
code/diff_sv/arguments.py
wngh1187/Diff-SV
import random
import math
import copy
import torch
import torch.utils.data as td
import soundfile as sf
import numpy as np

import utils.util as util
from data.noise_aug import Noises
from data.reverberation import RIRReverberation
from torch_audiomentations import Compose, Gain, AddColoredNoise, PitchShift

def get_loaders(args, vox1, voices):
    train_set = TrainSet(args, vox1.train_set)
    test_set = TestSet(args, vox1.test_set)
    dev_set_voices = TestSet_voices(args, voices.dev_set)
    eval_set_voices = TestSet_voices(args, voices.eval_set)
    if len(args['gpu_ids']) >1:
        train_set_sampler = td.DistributedSampler(train_set, shuffle=True)    
        test_set_sampler = td.DistributedSampler(test_set, shuffle=False)
        dev_set_sampler_voices = td.DistributedSampler(dev_set_voices, shuffle=False)
        eval_set_sampler_voices = td.DistributedSampler(eval_set_voices, shuffle=False)
    else:
        train_set_sampler = None
        test_set_sampler = None
        dev_set_sampler_voices = None
        eval_set_sampler_voices = None
    

    train_loader = td.DataLoader(
        train_set,
        batch_size=args['batch_size'],
        pin_memory=True,
        num_workers=args['num_workers'],
        sampler=train_set_sampler,
        drop_last=True,
        shuffle=False if train_set_sampler is not None else True
    )

    test_loader = td.DataLoader(
        test_set,
        batch_size=args['batch_size']//4,
        pin_memory=True,
        num_workers=args['num_workers']//2,
        sampler=test_set_sampler,
        shuffle = False
    )

    dev_loader_voices = td.DataLoader(
        dev_set_voices,
        batch_size=args['batch_size']//4,
        pin_memory=True,
        num_workers=args['num_workers']//2,
        sampler=dev_set_sampler_voices,
        shuffle = False
    )

    eval_loader_voices = td.DataLoader(
        eval_set_voices,
        batch_size=args['batch_size']//4,
        pin_memory=True,
        num_workers=args['num_workers']//2,
        sampler=eval_set_sampler_voices,
        shuffle = False
    )
    
    return train_set_sampler, train_loader, test_set, test_loader, dev_loader_voices, eval_loader_voices

class TrainSet(td.Dataset):
    def __init__(self, args, dataset):
        self.items = dataset
        self.args = args
        
        # set label
        count = 0

        # crop size
        self.crop_size = args['winlen'] + (args['winstep'] * (args['train_frame'] - 1))

        self.p = args['DA_p']
    
        # musan
        self.noises = Noises(f'{args["path_musan"]}_split/train')
        
        # rir
        self.rir = RIRReverberation(args['path_rir'])

        # aug
        self.apply_augmentation = Compose(
            transforms=[
                Gain(
                    min_gain_in_db=-15.0,
                    max_gain_in_db=5.0,
                    p=0.2,
                ),
                PitchShift(p=0.2, sample_rate=16000),
                AddColoredNoise(p=0.2)
            ]
            )

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):

        item = self.items[index]

        # read wav
        audio, _ = sf.read(item.path)
        
        # rand crop
        audio = torch.from_numpy(audio).float()
        audio = util.rand_crop(audio, self.crop_size)

        if self.p < random.random():
            return audio, audio, item.label, False
        else:
            ref_audio = copy.deepcopy(audio)
            if len(self.args['gpu_ids']) >1:
                audio = self.apply_augmentation(audio.unsqueeze(0).unsqueeze(0).to(self.args['device']), sample_rate=16000).squeeze(0).squeeze(0).cpu()
            else:
                audio = self.apply_augmentation(audio.unsqueeze(0).unsqueeze(0), sample_rate=16000).squeeze(0).squeeze(0)
            aug_type = random.randint(0, 4)
            if aug_type == 0:
                audio = self.rir(audio)
            elif aug_type == 1:
                audio = self.noises(audio, ['speech'])
            elif aug_type == 2:
                audio = self.noises(audio, ['music'])
            elif aug_type == 3:
                audio = self.noises(audio, ['noise'])
            elif aug_type == 4:
                audio = self.noises(audio, ['speech', 'music'])

            return audio, ref_audio, item.label, True
               
        
class TestSet(td.Dataset):
    @property
    def Key(self):
        return self.key
    @Key.setter
    def Key(self, value):
        self.key = value

    def __init__(self, args, dataset):
        self.key = 'clean'
        self.items = dataset
        self.crop_size = args['winlen'] + (args['winstep'] * (args['test_frame'] - 1))
        self.num_seg = args['num_seg']
        
    def __len__(self):
        return len(self.items[self.Key])

    def __getitem__(self, index):
        item = self.items[self.Key][index]

        # read wav
        audio, _ = sf.read(item.path)

        # crop
        if audio.shape[0] <= self.crop_size:
            shortage = self.crop_size - audio.shape[0]
            audio = np.pad(audio, (0, shortage), 'wrap')

        # stack
        buffer = []
        indices = np.linspace(0, audio.shape[0] - self.crop_size, self.num_seg)
        for idx in indices:
            idx = int(idx)
            buffer.append(audio[idx:idx + self.crop_size])
        buffer = np.stack(buffer, axis=0)

        return buffer.astype(np.float), item.key

class TestSet_voices(td.Dataset):
    @property
    def Key(self):
        return self.key
    @Key.setter
    def Key(self, value):
        self.key = value

    def __init__(self, args, dataset):
        self.items = dataset
        self.crop_size = args['winlen'] + (args['winstep'] * (args['test_frame'] - 1))
        self.num_seg = args['num_seg']
        
    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        item = self.items[index]

        # read wav
        audio, _ = sf.read(item.path)

        # crop
        if audio.shape[0] <= self.crop_size:
            shortage = self.crop_size - audio.shape[0]
            audio = np.pad(audio, (0, shortage), 'wrap')

        # stack
        buffer = []
        indices = np.linspace(0, audio.shape[0] - self.crop_size, self.num_seg)
        for idx in indices:
            idx = int(idx)
            buffer.append(audio[idx:idx + self.crop_size])
        buffer = np.stack(buffer, axis=0)

        return buffer.astype(np.float), item.key

def round_down(num, divisor):
    return num - (num % divisor)
code/diff_sv/data/data_loader.py
wngh1187/Diff-SV
import os
import random
import numpy as np
import soundfile as sf
import torch 

class Noises:
    Category = ['noise','speech','music', 'nonspeech']
    
    SNR = {
        'noise': (0, 20),
        'speech': (0, 20),
        'music': (0, 20),
        'nonspeech': (0, 20)
    }
    
    NumFile = {
        'noise': (1, 1),
        'speech': (3, 6),
        'music': (1, 1),
        'nonspeech': (1, 1)
    }
    
    def __init__(self, musan_path, nonspeech_path=None):
        # set vars
        self.noise_list = {}
        self.num_noise_file = {}
        for category in self.Category:
            self.noise_list[category] = []
            self.num_noise_file[category] = 0

        # init noise list
        for root, _, files in os.walk(musan_path):
            if self.Category[0] in root:
                category = self.Category[0]
            elif self.Category[1] in root:
                category = self.Category[1]
            elif self.Category[2] in root:
                category = self.Category[2]

            for file in files:
                if '.wav' in file and '._' not in file:
                    self.noise_list[category].append(os.path.join(root, file))
                    self.num_noise_file[category] += 1
        if nonspeech_path:
            for root, _, files in os.walk(nonspeech_path):
                for file in files:
                    if '.wav' in file and '._' not in file:
                        self.noise_list['nonspeech'].append(os.path.join(root, file))
                        self.num_noise_file['nonspeech'] += 1

    def __call__(self, x, categories):
        # calculate dB
        x = x.numpy() if type(x) == torch.Tensor else x
        x_size = x.shape[-1]
        
        x_dB = self.calculate_decibel(x)

        l_files = []

        if categories == 'none': return x
        for category in categories:
            snr_min, snr_max = self.SNR[category]
            file_min, file_max = self.NumFile[category]
            l_files.extend(random.sample(
                self.noise_list[category],
                random.randint(file_min, file_max)
            ))

        # init noise
        noises = []
        for f in l_files:
            # read .wav
            noise, _ = sf.read(f)
            # random crop
            noise_size = noise.shape[0]
            if noise_size < x_size:
                shortage = x_size - noise_size + 1
                noise = np.pad(
                    noise, (0, shortage), 'wrap'
                )
                noise_size = noise.shape[0]
            index = random.randint(0, noise_size - x_size)

            noises.append(noise[index:index + x_size])

        # noise injection
        noise = np.mean(noises, axis=0)
        # set SNR
        snr = random.uniform(snr_min, snr_max)
        # calculate dB
        noise_dB = self.calculate_decibel(noise)
        # append
        p = (x_dB - noise_dB - snr)
        
        noise = np.sqrt(10 ** (p / 10)) * noise
        x += noise
        
        return torch.from_numpy(x)

    def calculate_decibel(self, wav):
        assert 0 <= np.mean(wav ** 2) + 1e-4 
        return 10 * np.log10(np.mean(wav ** 2) + 1e-4)

code/diff_sv/data/noise_aug.py
wngh1187/Diff-SV
import os
import soundfile as sf
from tqdm import tqdm
from scipy.io import wavfile
from .noise_aug import Noises

class DataPreprocessor:
    def __init__(self, path_musan, path_nonspeech, path_vox1_test):
        self.path_musan = path_musan
        self.path_nonspeech = path_nonspeech
        self.path_vox1_test = path_vox1_test

        dir = self._leaf_dir(self.path_musan)
        self.path_musan_split = self.path_musan.replace(dir, f'{dir}_split')

        self.dir_vox1_test = self._leaf_dir(self.path_vox1_test)
        self.path_vox1_test_noise = self.path_vox1_test.replace(
                self.dir_vox1_test, f'{self.dir_vox1_test}_noise')

    def _leaf_dir(self, path):
        """
        Return leaf sub directory name from path
        """
        if '.' in path:
            name = os.path.dirname(path).split('/')[-1]
        elif path[-1] == '/':
            name = os.path.dirname(path).split('/')[-1]
        else:
            name = os.path.basename(path)

        return name

    def check_environment(self):
        """Check if preprocessing has been done
        """
        if not os.path.exists(self.path_musan_split):
            print('You need to split musan dataset')
            print('Now processing...')
            self.split_musan()
        
        if not os.path.exists(self.path_vox1_test_noise):		
            self.noises = Noises(self.path_musan_split + '/test', self.path_nonspeech)
            print('You need to make noisy test set')
            print('Now processing...')
            self.init_noisey_test_set()

    def split_musan(self):
        """
        Split MUSAN dataset for 
            - to divide train/test set
            - to increase reading speed
        """
        winlen = 16000 * 5
        winstep = 16000 * 3

        for root, _, files in os.walk(self.path_musan):
            for file in tqdm(files):
                if '.wav' in file:
                    num = int(file[-8:-4])
                    assert 0 <= num and num < 1000
                    category = 'train' if num % 2 == 0 else 'test'

                    file = os.path.join(root, file)
                    temp = f'{self.path_musan_split}/{category}'
                    destination = file.replace(self.path_musan, temp
                                    ).replace('.', '')
                    os.makedirs(destination, exist_ok=True)

                    sr, data = wavfile.read(file)
                    num_file = (len(data) - winlen) // winstep

                    if 0 < num_file:
                        for i in range(num_file):
                            wavfile.write(f'{destination}/{i * 3}_{(i * 3) + 5}.wav', sr, data[i * winstep:i * winstep + winlen])
                    else:
                        wavfile.write(f'{destination}/all.wav', sr, data)

    def init_noisey_test_set(self):
        """Inject musan noise to vox1 testset for making noisy test set
        """
        for root, _, files in tqdm(os.walk(self.path_vox1_test)):
            for file in files:
                if '.wav' in file:
                    path = os.path.join(root, file)
                    self._inject_noise(path, 'noise', 0)
                    self._inject_noise(path, 'noise', 5)
                    self._inject_noise(path, 'noise', 10)
                    self._inject_noise(path, 'noise', 15)
                    self._inject_noise(path, 'noise', 20)
                    self._inject_noise(path, 'speech', 0)
                    self._inject_noise(path, 'speech', 5)
                    self._inject_noise(path, 'speech', 10)
                    self._inject_noise(path, 'speech', 15)
                    self._inject_noise(path, 'speech', 20)
                    self._inject_noise(path, 'music', 0)
                    self._inject_noise(path, 'music', 5)
                    self._inject_noise(path, 'music', 10)
                    self._inject_noise(path, 'music', 15)
                    self._inject_noise(path, 'music', 20)
                    self._inject_noise(path, 'nonspeech', 0)
                    self._inject_noise(path, 'nonspeech', 5)
                    self._inject_noise(path, 'nonspeech', 10)
                    self._inject_noise(path, 'nonspeech', 15)
                    self._inject_noise(path, 'nonspeech', 20)

    def _inject_noise(self, path, category, snr):
        data, sr = sf.read(path)
        self.noises.Category = [category]
        self.noises.SNR[category] = (snr, snr)
        destination = path.replace(
            self.dir_vox1_test, f'{self.dir_vox1_test}_noise/{category}_{snr}')
        os.makedirs(os.path.dirname(destination), exist_ok=True)
        data = self.noises(data, [category])
        sf.write(destination, data, sr)
code/diff_sv/data/preprocessing.py
wngh1187/Diff-SV
import os
import random
import numpy as np
import soundfile as sf
from scipy import signal
import torch

class RIRReverberation:
    def __init__(self, path):
        self.files = []
        
        # parse list
        for root, _, files in os.walk(path):
            for file in files:
                if '.wav' in file:
                    self.files.append(os.path.join(root, file))


    def __call__(self, x):
        path = random.sample(self.files, 1)[0]
        
        rir, _ = sf.read(path)
        rir = rir.astype(np.float)
        rir = np.expand_dims(rir, 0)
        rir = rir / np.sqrt(np.sum(rir**2))
        
        x = np.expand_dims(x, 0)
        x = signal.convolve(x, rir, mode='full')[:,:len(x[0])]

        x = np.squeeze(x, 0)

        return torch.from_numpy(x)

code/diff_sv/data/reverberation.py
wngh1187/Diff-SV
import os
from dataclasses import dataclass

@dataclass
class TestItem:
    key: str
    path: str

@dataclass
class TestTrial:
    key1: str
    key2: str
    label: str

class Voices:
    @property
    def dev_set(self):
        return self.__dev_set

    @property
    def eval_set(self):
        return self.__eval_set

    @property
    def dev_trials(self):
        return self.__dev_trials

    @property
    def eval_trials(self):
        return self.__eval_trials

    def __init__(self, path_dev, path_dev_trial, path_eval, path_eval_trial):
        
        # dev_set
        self.__dev_set = []
        self._parse_item(path_dev, self.__dev_set)

        # dev_set
        self.__eval_set = []
        self._parse_item(path_eval, self.__eval_set)
        
        # test_trials
        self.__dev_trials = self._parse_trials(path_dev_trial)
        
        # test_trials
        self.__eval_trials = self._parse_trials(path_eval_trial)
        
    def _parse_item(self, path, dataset):    
        for root, _, files in os.walk(path):
            for file in files:
                # if '.wav' in file:
                if file[-4:] == '.wav'  and file[:2] != '._':
                    temp = os.path.join(root, file)
                    dataset.append(
                        TestItem(
                            path=temp,
                            key=file[:-4]
                        )
                    )

    def _parse_trials(self, path):
        trials = []

        f = open(path) 
        for line in f.readlines():
            strI = line.split(' ')
            trials.append(
                TestTrial(
                    key1=strI[0].replace('\n', ''),
                    key2=strI[1].split('/')[-1].split('.')[0].replace('\n', ''),
                    label='0' if strI[-1].strip() == 'imp' else '1'
                )
            )
        
        return trials
code/diff_sv/data/voices.py
wngh1187/Diff-SV
import os
from dataclasses import dataclass

NUM_TRAIN_ITEM = 148642
NUM_TRAIN_SPK = 1211
NUM_TRIALS = 37611
NUM_TEST_ITEM = 4874

@dataclass
class TrainItem:
    path: str
    speaker: str
    label: int

@dataclass
class TestItem:
    key: str
    path: str

@dataclass
class TestTrial:
    key1: str
    key2: str
    label: str

class VoxCeleb1:
    def __init__(self, path_train, path_test, path_test_noise, path_trial):

        self.train_set = []

        # train_set
        labels = {}
        num_utt = [0 for _ in range(NUM_TRAIN_SPK)]
        num_sample = 0
        for root, _, files in os.walk(path_train):
            for file in files:
                if '.wav' in file:
                    # combine path
                    f = os.path.join(root, file)
                    
                    # parse speaker
                    spk = f.split('/')[-3]
                    
                    # labeling
                    try: labels[spk]
                    except: 
                        labels[spk] = len(labels.keys())

                    # init item
                    item = TrainItem(path=f, speaker=spk, label=labels[spk])
                    self.train_set.append(item)
                    num_sample += 1
                    num_utt[labels[spk]] += 1         

        # test_set
        self.test_set = {
            'clean': [],
            'noise_0': [],
            'noise_5': [],
            'noise_10': [],
            'noise_15': [],
            'noise_20': [],
            'speech_0': [],
            'speech_5': [],
            'speech_10': [],
            'speech_15': [],
            'speech_20': [],
            'music_0': [],
            'music_5': [],
            'music_10': [],
            'music_15': [],
            'music_20': [],
            'nonspeech_0': [],
            'nonspeech_5': [],
            'nonspeech_10': [],
            'nonspeech_15': [],
            'nonspeech_20': []
        }
        self._parse_item(path_test, 'clean')
        self._parse_item(f'{path_test_noise}/noise_0', 'noise_0')
        self._parse_item(f'{path_test_noise}/noise_5', 'noise_5')
        self._parse_item(f'{path_test_noise}/noise_10', 'noise_10')
        self._parse_item(f'{path_test_noise}/noise_15', 'noise_15')
        self._parse_item(f'{path_test_noise}/noise_20', 'noise_20')

        self._parse_item(f'{path_test_noise}/speech_0', 'speech_0')
        self._parse_item(f'{path_test_noise}/speech_5', 'speech_5')
        self._parse_item(f'{path_test_noise}/speech_10', 'speech_10')
        self._parse_item(f'{path_test_noise}/speech_15', 'speech_15')
        self._parse_item(f'{path_test_noise}/speech_20', 'speech_20')

        self._parse_item(f'{path_test_noise}/music_0', 'music_0')
        self._parse_item(f'{path_test_noise}/music_5', 'music_5')
        self._parse_item(f'{path_test_noise}/music_10', 'music_10')
        self._parse_item(f'{path_test_noise}/music_15', 'music_15')
        self._parse_item(f'{path_test_noise}/music_20', 'music_20')

        self._parse_item(f'{path_test_noise}/nonspeech_0', 'nonspeech_0')
        self._parse_item(f'{path_test_noise}/nonspeech_5', 'nonspeech_5')
        self._parse_item(f'{path_test_noise}/nonspeech_10', 'nonspeech_10')
        self._parse_item(f'{path_test_noise}/nonspeech_15', 'nonspeech_15')
        self._parse_item(f'{path_test_noise}/nonspeech_20', 'nonspeech_20')

        # test_trials
        self.test_trials = self._parse_trials(path_trial)
        
        # error check
        assert len(self.train_set) == 148642, f'len(train_set): {len(self.train_set)}'
        assert len(self.test_set["clean"]) == 4874, f'len(test_set): {len(self.test_set["clean"])}'
        assert len(self.test_set["noise_0"]) == 4874, f'len(test_set): {len(self.test_set["noise_0"])}'
        assert len(self.test_set["noise_5"]) == 4874, f'len(test_set): {len(self.test_set["noise_5"])}'
        assert len(self.test_set["noise_10"]) == 4874, f'len(test_set): {len(self.test_set["noise_10"])}'
        assert len(self.test_set["noise_15"]) == 4874, f'len(test_set): {len(self.test_set["noise_15"])}'
        assert len(self.test_set["noise_20"]) == 4874, f'len(test_set): {len(self.test_set["noise_20"])}'
        assert len(self.test_set["speech_0"]) == 4874, f'len(test_set): {len(self.test_set["speech_0"])}'
        assert len(self.test_set["speech_5"]) == 4874, f'len(test_set): {len(self.test_set["speech_5"])}'
        assert len(self.test_set["speech_10"]) == 4874, f'len(test_set): {len(self.test_set["speech_10"])}'
        assert len(self.test_set["speech_15"]) == 4874, f'len(test_set): {len(self.test_set["speech_15"])}'
        assert len(self.test_set["speech_20"]) == 4874, f'len(test_set): {len(self.test_set["speech_20"])}'
        assert len(self.test_set["music_0"]) == 4874, f'len(test_set): {len(self.test_set["music_0"])}'
        assert len(self.test_set["music_5"]) == 4874, f'len(test_set): {len(self.test_set["music_5"])}'
        assert len(self.test_set["music_10"]) == 4874, f'len(test_set): {len(self.test_set["music_10"])}'
        assert len(self.test_set["music_15"]) == 4874, f'len(test_set): {len(self.test_set["music_15"])}'
        assert len(self.test_set["music_20"]) == 4874, f'len(test_set): {len(self.test_set["music_20"])}'
        assert len(self.test_set["nonspeech_0"]) == 4874, f'len(test_set): {len(self.test_set["nonspeech_0"])}'
        assert len(self.test_set["nonspeech_5"]) == 4874, f'len(test_set): {len(self.test_set["nonspeech_5"])}'
        assert len(self.test_set["nonspeech_10"]) == 4874, f'len(test_set): {len(self.test_set["nonspeech_10"])}'
        assert len(self.test_set["nonspeech_15"]) == 4874, f'len(test_set): {len(self.test_set["nonspeech_15"])}'
        assert len(self.test_set["nonspeech_20"]) == 4874, f'len(test_set): {len(self.test_set["nonspeech_20"])}'
        assert len(self.test_trials) == 37720, f'len(test_trials): {len(self.test_trials)}'

    def _parse_item(self, path, key):    
        for root, _, files in os.walk(path):
            for file in files:
                if '.wav' in file and len(file) == 9:                    
                    temp = os.path.join(root, file)
                    self.test_set[key].append(
                        TestItem(
                            path=temp,
                            key='/'.join(temp.split('/')[-3:])
                        )
                    )

    def _parse_trials(self, path):
        trials = []

        f = open(path) 
        for line in f.readlines():
            strI = line.split(' ')
            trials.append(
                TestTrial(
                    key1=strI[1].replace('\n', ''),
                    key2=strI[2].replace('\n', ''),
                    label=strI[0]
                )
            )
        
        return trials

code/diff_sv/data/voxceleb1.py
wngh1187/Diff-SV
import os
from .interface import ExperimentLogger
from .local import LocalLogger
from .wandb import WandbLogger

class LogModuleController(ExperimentLogger):
    def __init__(self, modules):
        self.components = modules

    def log_metric(self, name, value, log_step=None, epoch_step=None):
        for component in self.components:
            component.log_metric(name, value, log_step, epoch_step)

    def log_text(self, name, text):
        for component in self.components:
            component.log_text(name, text)

    def log_parameter(self, dictionary):
        for component in self.components:
            component.log_parameter(dictionary)

    def log_image(self, name, image):
        for component in self.components:
            component.log_image(name, image)

    def save_model(self, name, model):
        for component in self.components:
            component.save_model(name, model)
            
    def finish(self):
        for component in self.components:
            component.finish()
    
    class Builder():
        def __init__(self, name, project):
            self.args = {}
            self.args['name'] = name
            self.args['project'] = project
            self.args['flag_use_local'] = False
            self.args['flag_use_wandb'] = False
            self.args['tags'] = []
            self.args['description'] = ''
            self.args['source_path'] = None
            self.args['source_files'] = None
        
        def tags(self, tags):
            if type(tags) is list:
                self.args['tags'] = tags
                return self
            else:
                raise Exception(f'tags must be list. Given type is {type(tags)}')

        def description(self, description):
            self.args['description'] = description
            return self

        def save_source_files(self, path):
            self.args['source_path'] = path
            self.args['source_files'] = []
            for p, _, fs in os.walk(path):
                for f in fs:
                    ext = os.path.splitext(f)[-1]
                    if ext == '.py' and 'vscode' not in p:
                        self.args['source_files'].append(
                            f"{p}/{f}")
            return self

        def use_local(self, path):
            self.args['flag_use_local'] = True
            self.args['local_path'] = path
            return self
        
        def use_wandb(self, group, entity):
            self.args['flag_use_wandb'] = True
            self.args['group'] = group
            self.args['entity'] = entity
            return self

        def build(self):
            modules = []

            if self.args['flag_use_local']:
                local = LocalLogger(
                    self.args['local_path'], 
                    self.args['name'], 
                    self.args['project'], 
                    self.args['tags'], 
                    self.args['description'],
                    self.args['source_path']
                )
                modules.append(local)

            if self.args['flag_use_wandb']:
                wandb = WandbLogger(
                    self.args['local_path'], 
                    self.args['name'],
                    self.args['group'],
                    self.args['project'],
                    self.args['entity'],
                    self.args['tags'],
                    self.args['source_path'],
                    os.path.join("/", self.args['local_path'], self.args['project'], self.args['tags'][0], self.args['name'])
                )
                modules.append(wandb)

            package = LogModuleController(modules)

            return package
code/diff_sv/log/controller.py
wngh1187/Diff-SV
from abc import ABCMeta, abstractmethod

class ExperimentLogger(metaclass=ABCMeta):
    @abstractmethod
    def log_metric(self, name, value, log_step=None, epoch_step=None):
        pass

    @abstractmethod
    def log_text(self, name, text):
        pass

    @abstractmethod
    def log_parameter(self, dictionary):
        pass

    @abstractmethod    
    def log_image(self, name, image):
        pass

    @abstractmethod
    def save_model(self, name, state_dict):
        pass

    @abstractmethod
    def finish(self):
        pass
code/diff_sv/log/interface.py
wngh1187/Diff-SV
import os
import time
import torch
import shutil
import zipfile
from threading import Thread
from .interface import ExperimentLogger

def zipdir(path, ziph):
    for root, dirs, files in os.walk(path):
        for file in files:
        
            fn, ext = os.path.splitext(file)
            if ext != ".py":
                continue
            
            ap = '/'.join(os.path.abspath(file).split('/')[:-1])
            ziph.write(os.path.join(root, file))

class LocalLogger(ExperimentLogger):
    """Save experiment logs to local storage.
    Note that logs are synchronized every 10 seconds.
    """
    def __init__(self, path, name, project, tags, description=None, script=None):
        # set path
        path = os.path.join("/", path, project, "/".join(tags), name)
        self.paths = {
            'metric': f'{path}/metric',
            'text'  : f'{path}/text',
            'image' : f'{path}/image',
            'model' : f'{path}/model',
            'script': f'{path}/script',
        }
        
        # make directory
        # if os.path.exists(path):
        # 	shutil.rmtree(path)
        for dir in self.paths.values():
            if not os.path.exists(dir): os.makedirs(dir) 

        # save description
        if description is not None:
            self.log_text('description', description)

        # zip codes
        with zipfile.ZipFile(
            self.paths['script'] + "/script.zip", "w", zipfile.ZIP_DEFLATED
        ) as f_zip:
            zipdir(script, f_zip)

        # synchronize thread
        self._metrics = {}
        thread_synchronize = Thread(target=self._sync)
        thread_synchronize.setDaemon(True)
        thread_synchronize.start()

    def _sync(self):
        """Flush files
        """
        while(True):
            for txt in self._metrics.values():
                txt.flush()
            time.sleep(10)

    def log_text(self, name, text):
        path = f'{self.paths["text"]}/{name}.txt'
        mode = 'a' if os.path.exists(path) else 'w'
        file = open(path, mode, encoding='utf-8')
        file.write(text)
        file.close()

    def log_parameter(self, dictionary):
        for k, v in dictionary.items():
            self.log_text('parameters', f'{k}: {v}\n')

    def log_metric(self, name, value, log_step=None, epoch_step=None):
        if name not in self._metrics.keys():
            path = f'{self.paths["metric"]}/{name}.txt'
            self._metrics[name] = open(path, 'w', encoding='utf-8').close()
            self._metrics[name] = open(path, 'a', encoding='utf-8')

        if log_step is None:
            self._metrics[name].write(f'{name}: {value}\n')
        else:
            self._metrics[name].write(f'[{log_step}] {name}: {value}\n')
    
    def log_image(self, name, image):
        path = f'{self.paths["image"]}/{name}.png'
        image.save(path, 'PNG')
    
    def save_model(self, name, state_dict):
        path = f'{self.paths["model"]}/{name}.pt'
        torch.save(state_dict, path)

    def finish(self):
        pass
code/diff_sv/log/local.py
wngh1187/Diff-SV
import os
import torch
import wandb
from wandb import AlertLevel
from .interface import ExperimentLogger

class WandbLogger(ExperimentLogger):
    """Save experiment logs to wandb
    """
    def __init__(self, path, name, group, project, entity, tags, script=None, save_dir = None):
        self.paths = {
            'model' : f'{save_dir}/model'
        }
        self.run = wandb.init(
                group=group,
                project=project,
                entity=entity,
                tags=tags,
                dir=save_dir
            )
        wandb.run.name = name
        # upload zip file
        wandb.save(save_dir + "/script/script.zip")


    def log_metric(self, name, value, log_step=None, epoch_step=None):
        if epoch_step is not None:
            wandb.log({
                name: value,
                'epoch_step': epoch_step})
        elif log_step is not None:
            wandb.log({
                name: value,
                'log_step': log_step})

    def log_text(self, name, text):
        pass

    def log_image(self, name, image):
        wandb.log({name: [wandb.Image(image)]})

    def log_parameter(self, dictionary):
        wandb.config.update(dictionary)

    def save_model(self, name, state_dict):
        path = f'{self.paths["model"]}/BestModel.pt'
        wandb.save(path)

    def finish(self):
        wandb.finish()
code/diff_sv/log/wandb.py
wngh1187/Diff-SV
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

class LossFunction(nn.Module):
    def __init__(self, nOut, nClasses, margin=0.2, scale=30, easy_margin=False, **kwargs):
        super(LossFunction, self).__init__()

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

        print('Initialised AAM_Softmax Loss')

    def forward(self, x, label=None):
        assert x.size()[0] == label.size()[0]
        assert x.size()[1] == self.in_feats

        # 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

        loss = self.ce(output, label)
        return loss

code/diff_sv/loss/aam_softmax.py
wngh1187/Diff-SV
import torch.nn as nn

class LossFunction(nn.Module):
	def __init__(self, nOut, nClasses, **kwargs):
		super(LossFunction, self).__init__()

		self.test_normalize = True
		
		self.criterion  = nn.CrossEntropyLoss()
		self.fc 		= nn.Linear(nOut, nClasses, bias=True)

		print('Initialised Softmax Loss')

	def forward(self, x, label=None):

		x 		= self.fc(x)
		nloss   = self.criterion(x, label)
		return nloss
code/diff_sv/loss/softmax.py
wngh1187/Diff-SV
import os
import random
import importlib
import numpy as np
import torch
import torch.nn as nn

import arguments
import trainers.train as train
import data.data_loader as data
from data.voxceleb1 import VoxCeleb1
from data.voices import Voices
from optimizer.CosineAnnealingWarmUpRestarts import CosineAnnealingWarmUpRestarts
from models.diffsv import DiffSV
from log.controller import LogModuleController
from data.preprocessing import DataPreprocessor

def set_experiment_environment(args):
    #### reproducible ###
    random.seed(args['rand_seed'])
    np.random.seed(args['rand_seed'])
    torch.manual_seed(args['rand_seed'])
    torch.backends.cudnn.deterministic = args['flag_reproduciable']
    torch.backends.cudnn.benchmark = not args['flag_reproduciable']

    ### DDP env ###
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '4021'
    args['rank'] = args['process_id']
    args['device'] = 'cuda:' + args['gpu_ids'][args['process_id']]
    torch.cuda.empty_cache()
    torch.cuda.set_device(args['device'])
    if len(args['gpu_ids']) >1:
        torch.distributed.init_process_group(
            backend='nccl', world_size=args['world_size'], rank=args['rank'])

def run(process_id, args):
    ### check parent process ###
    args['process_id'] = process_id
    args['flag_parent'] = process_id == 0
    
    ### experiment environment ###
    set_experiment_environment(args)
    trainer = train.ModelTrainer(args)
    trainer.args = args
    
    ### set logger ###
    if args['flag_parent']:
        logger = LogModuleController.Builder(args['name'], args['project']
        ).tags(args['tags']
        ).description(args['description']
        ).save_source_files(args['path_scripts']
        ).use_local(args['path_logging']
        #).use_wandb(args['wandb_group'], args['wandb_entity'] #Activate if using WandB (https://wandb.ai)
        ).build()
        trainer.logger = logger
    
    ### set dataset ###
    trainer.vox1 = VoxCeleb1(
        args['path_vox1_train'], 
        args['path_vox1_test'], 
        f'{args["path_vox1_test"]}_noise', 
        args['path_vox1_trials']
    )
    args['num_speaker'] = 1211

    trainer.voices = Voices(
        args['path_voices_dev'],
        args['path_voices_dev_trials'],
        args['path_voices_eval'],
        args['path_voices_eval_trials']
    )

    ### set data loader ###
    loaders = data.get_loaders(args, trainer.vox1, trainer.voices)
    trainer.train_set_sampler = loaders[0]
    trainer.train_loader = loaders[1]
    trainer.enrollment_set = loaders[2]
    trainer.enrollment_loader = loaders[3]
    trainer.dev_loader_voices = loaders[4]
    trainer.eval_loader_voices = loaders[5]

    ### set model ###
    model = DiffSV(args).to(args['device'])

    if args['flag_parent']:
        args['model_num_param'] = model.nparams
        print('num of model params: ', args['model_num_param'])
        trainer.logger.log_parameter(args)

    ### set criterion ###
    criterion = {}
    classification_loss_function = importlib.import_module('loss.'+ args['classification_loss']).__getattribute__('LossFunction')
    criterion['classification_loss'] = classification_loss_function(args['code_dim'], args['num_speaker']).to(args['device']) 
    criterion_params = list(criterion['classification_loss'].parameters())

    ### set ddp if using multi gpu ###
    if len(args['gpu_ids']) >1:   
        model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
        model = nn.parallel.DistributedDataParallel(model, device_ids=[args['device']], find_unused_parameters=True)
        criterion['classification_loss'] = nn.parallel.DistributedDataParallel(criterion['classification_loss'], device_ids=[args['device']], find_unused_parameters=True)

    trainer.model = model
    trainer.criterion = criterion

    ### set optimizer ###
    trainer.optimizer = torch.optim.Adam(
        list(model.parameters()) + criterion_params, 
        lr=args['lr_start'], 
        weight_decay=args['weigth_decay'],
        amsgrad = args['amsgrad']
    )

    ### set lr scheduler ###
    args['number_iteration'] = len(trainer.train_loader)
    trainer.lr_scheduler = CosineAnnealingWarmUpRestarts(
        trainer.optimizer, 
        T_0=args['number_iteration'] * args['number_cycle'], 
        T_mult=args['T_mult'], 
        eta_max=args['eta_max'],  
        T_up=args['number_iteration'] * args['warmup_number_cycle'], 
        gamma=args['gamma'])

    ### train and test model ###
    trainer.run()
        
    if args['flag_parent']:    trainer.logger.finish()


if __name__ == '__main__':
    ### get arguments ###
    args = arguments.get_args()

    ### check dataset ###
    data_preprocessor = DataPreprocessor(args['path_musan'], args['path_nonspeech'], args['path_vox1_test'])
    data_preprocessor.check_environment()

    ### set gpu device ###
    if args['usable_gpu'] is None: 
        args['gpu_ids'] = os.environ['CUDA_VISIBLE_DEVICES'].split(',')
    else:
        args['gpu_ids'] = args['usable_gpu'].split(',')
    
    if len(args['gpu_ids']) == 0 or len(args['gpu_ids']) == 1:
        run(0, args)
    else:
        ### set DDP ###
        args['world_size'] = len(args['gpu_ids'])
        args['batch_size'] = args['batch_size'] // (args['world_size'])
        args['num_workers'] = args['num_workers'] // args['world_size']
        
        torch.multiprocessing.set_sharing_strategy('file_system')
        torch.multiprocessing.spawn(
            run, 
            nprocs=args['world_size'], 
            args=(args,)
        )
code/diff_sv/main.py
wngh1187/Diff-SV
import numpy as np
import torch


class BaseModule(torch.nn.Module):
    def __init__(self):
        super(BaseModule, self).__init__()

    @property
    def nparams(self):
        """
        Returns number of trainable parameters of the module.
        """
        num_params = 0
        for name, param in self.named_parameters():
            if param.requires_grad:
                num_params += np.prod(param.detach().cpu().numpy().shape)
        return num_params


    def relocate_input(self, x: list):
        """
        Relocates provided tensors to the same device set for the module.
        """
        device = next(self.parameters()).device
        for i in range(len(x)):
            if isinstance(x[i], torch.Tensor) and x[i].device != device:
                x[i] = x[i].to(device)
        return x

code/diff_sv/models/base.py
wngh1187/Diff-SV
import torch
from torch import nn
import torch.nn.functional as F
from models.modules import GradLogPEstimator2d
from models.base import BaseModule


def get_noise(t, beta_init, beta_term, cumulative=False):
    if cumulative:
        noise = beta_init*t + 0.5*(beta_term - beta_init)*(t**2)
    else:
        noise = beta_init + (beta_term - beta_init)*t
    return noise

class Denoiser(BaseModule):
    def __init__(
            self,
            feature_channels,
            inference_steps=10,
    ):
        super().__init__()
        self.diffusion_inference_steps = inference_steps
        self.beta_min = 0.05
        self.beta_max = 20
        self.model = GradLogPEstimator2d(16, n_feats = feature_channels)
      
    def forward(self, x_hat, ref_x=None, is_inference=False):
        with torch.no_grad():
            z = x_hat + torch.randn_like(x_hat, device=x_hat.device) / 1.0
            z0_hat = self.inference(z, x_hat, 10)
        if is_inference: return z0_hat

        ddim_loss = self.calculate_loss(ref_x, x_hat)
        return z0_hat, ddim_loss

    def calculate_loss(self, x0, mu, offset=1e-5):
        t = torch.rand(x0.shape[0], dtype=x0.dtype, device=x0.device,
                       requires_grad=False)
        t = torch.clamp(t, offset, 1.0 - offset)

        time = t.unsqueeze(-1).unsqueeze(-1)
        cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True) #
        mean = x0*torch.exp(-0.5*cum_noise) + mu*(1.0 - torch.exp(-0.5*cum_noise))
        variance = 1.0 - torch.exp(-cum_noise)
        z = torch.randn(x0.shape, dtype=x0.dtype, device=x0.device, 
                        requires_grad=False)
        xt = mean + z * torch.sqrt(variance)
        
        time = t.unsqueeze(-1).unsqueeze(-1)
        cum_noise = get_noise(time, self.beta_min, self.beta_max, cumulative=True)

        noise_estimation = self.model(xt, mu, t)

        noise_estimation *= torch.sqrt(1.0 - torch.exp(-cum_noise))
        loss = torch.mean((noise_estimation + z)**2)
        return loss
    
    def inference(self, z, mu, n_timesteps):
        h = 1.0 / n_timesteps
        xt = z

        for i in range(n_timesteps):
            t = (1.0 - (i + 0.5)*h) * torch.ones(z.shape[0], dtype=z.dtype, 
                                                 device=z.device)
            time = t.unsqueeze(-1).unsqueeze(-1)
            noise_t = get_noise(time, self.beta_min, self.beta_max, 
                                cumulative=False)
            dxt = 0.5 * (mu - xt - self.model(xt, mu, t))
            dxt = dxt * noise_t * h
            xt = xt - dxt

        return xt


code/diff_sv/models/denoiser.py
wngh1187/Diff-SV
import torch
import torch.nn as nn
import numpy as np
from models.modules import SEBasicBlock
from models.denoiser import Denoiser as Denoiser
from models.transformer import Transformer
import torch.nn.functional as F
import torchaudio

from models.base import BaseModule
import math

class PreEmphasis(torch.nn.Module):

    def __init__(self, coef: float = 0.97):
        super().__init__()
        self.coef = coef
        self.register_buffer(
            'flipped_filter', torch.FloatTensor([-self.coef, 1.]).unsqueeze(0).unsqueeze(0)
        )

    def forward(self, input: torch.tensor) -> torch.tensor:
        input = input.unsqueeze(1)
        input = F.pad(input, (1, 0), 'reflect')
        return F.conv1d(input, self.flipped_filter).squeeze(1)


class DiffSV(BaseModule): 
    def __init__(self, args):
        super(DiffSV, self).__init__()
        self.max_epoch = args['epoch']
        self.l_channel = args['l_channel']
        self.l_num_convblocks = args['l_num_convblocks']
        self.code_dim = args['code_dim']
        self.stride = args['stride']
        self.first_kernel_size = args['first_kernel_size']
        self.first_stride_size = args['first_stride_size']
        self.first_padding_size = args['first_padding_size']

        self.torchfbank = nn.Sequential(
            PreEmphasis(),            
            torchaudio.transforms.MelSpectrogram(sample_rate=args['samplerate'], n_fft=args['nfft'], win_length=args['winlen'], hop_length=args['winstep'], \
                                                 f_min = args['f_min'], f_max = args['f_max'], window_fn=args['winfunc'], n_mels=args['nfilts']),
            )
        
        self.inplanes   = self.l_channel[0]
        
        ### define enhancer ###
        self.enhancer = Transformer(4, args['nfilts'], d_inner=512)

        #### define denoiser ###
        self.denoiser = Denoiser(args['nfilts'])

        ### define extractor ###
        self.conv1 = nn.Conv2d(3, self.l_channel[0] , kernel_size=self.first_kernel_size, stride=self.first_stride_size, padding=self.first_padding_size,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(self.l_channel[0])
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = self._make_layer(SEBasicBlock, self.l_channel[0], self.l_num_convblocks[0], stride=self.stride[0])
        self.layer2 = self._make_layer(SEBasicBlock, self.l_channel[1], self.l_num_convblocks[1], stride=self.stride[1])
        self.layer3 = self._make_layer(SEBasicBlock, self.l_channel[2], self.l_num_convblocks[2], stride=self.stride[2])
        self.layer4 = self._make_layer(SEBasicBlock, self.l_channel[3], self.l_num_convblocks[3], stride=self.stride[3])

        final_dim = self.l_channel[-1] * args['nfilts']//8
        
        self.attention = nn.Sequential(
            nn.Conv1d(final_dim, final_dim//8, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(final_dim//8),
            nn.Conv1d(final_dim//8, final_dim, kernel_size=1), 
            nn.Softmax(dim=-1),
        )
        
        self.bn_agg = nn.BatchNorm1d(final_dim * 2)

        self.fc = nn.Linear(final_dim*2, self.code_dim)
        self.bn_code = nn.BatchNorm1d(self.code_dim)
        
    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

        
    def forward(self, x, ref_x=None, aug=None, epoch=None, is_train=True):
        if is_train:
            with torch.no_grad():
                x = self.torchfbank(x)+1e-6
                x = x.log()   
                x = x - torch.mean(x, dim=-1, keepdim=True)
                org_x = x

                ref_x = self.torchfbank(ref_x)+1e-6
                ref_x = ref_x.log()   
                ref_x = ref_x - torch.mean(ref_x, dim=-1, keepdim=True)

            x_hat = self.enhancer(org_x)
            prior_loss = F.mse_loss(ref_x, x_hat)

            z0_hat, ddim_loss = self.denoiser(x_hat, ref_x)

            z0_hat = z0_hat.detach()
        else:
            with torch.no_grad():
                x = self.torchfbank(x)+1e-6
                x = x.log()   
                x = x - torch.mean(x, dim=-1, keepdim=True)
                org_x = x

            x_hat = self.enhancer(org_x)
            z0_hat = self.denoiser(x_hat, is_inference=True)

        x = torch.stack([org_x, x_hat, z0_hat], 1)

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        bs, _, _, time = x.size()
        x = x.reshape(bs, -1, time)
        
        w = self.attention(x)
        m = torch.sum(x * w, dim=-1)
        s = torch.sqrt((torch.sum((x ** 2) * w, dim=-1) - m ** 2).clamp(min=1e-5))
        x = torch.cat([m, s], dim=1)
        x =    self.bn_agg(x)

        code = self.fc(x)
        code = self.bn_code(code)

        
        if is_train:
            return code, prior_loss, ddim_loss
        else:
            return code
code/diff_sv/models/diffsv.py
wngh1187/Diff-SV
import torch
import torch.nn as nn
from torch.nn import functional as F
import numpy as np
from models.base import BaseModule
import math
from einops import rearrange

class Mish(BaseModule):
    def forward(self, x):
        return x * torch.tanh(torch.nn.functional.softplus(x))


class Upsample(BaseModule):
    def __init__(self, dim):
        super(Upsample, self).__init__()
        self.conv = torch.nn.ConvTranspose2d(dim, dim, 4, 2, 1)

    def forward(self, x):
        return self.conv(x)


class Downsample(BaseModule):
    def __init__(self, dim):
        super(Downsample, self).__init__()
        self.conv = torch.nn.Conv2d(dim, dim, 3, 2, 1)

    def forward(self, x):
        return self.conv(x)


class Rezero(BaseModule):
    def __init__(self, fn):
        super(Rezero, self).__init__()
        self.fn = fn
        self.g = torch.nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return self.fn(x) * self.g

class Block(BaseModule):
    def __init__(self, dim, dim_out, groups=8):
        super(Block, self).__init__()
        self.block = torch.nn.Sequential(torch.nn.Conv2d(dim, dim_out, 3, 
                                         padding=1), torch.nn.GroupNorm(
                                         groups, dim_out), Mish())
    def forward(self, x):
        output = self.block(x)
        return output

class ResnetBlock(BaseModule):
    def __init__(self, dim, dim_out, time_emb_dim, groups=8):
        super(ResnetBlock, self).__init__()
        self.mlp = torch.nn.Sequential(Mish(), torch.nn.Linear(time_emb_dim, 
                                                               dim_out))

        self.block1 = Block(dim, dim_out, groups=groups)
        self.block2 = Block(dim_out, dim_out, groups=groups)
        if dim != dim_out:
            self.res_conv = torch.nn.Conv2d(dim, dim_out, 1)
        else:
            self.res_conv = torch.nn.Identity()

    def forward(self, x, time_emb):
        h = self.block1(x)
        h += self.mlp(time_emb).unsqueeze(-1).unsqueeze(-1)
        h = self.block2(h)
        output = h + self.res_conv(x)
        return output


class LinearAttention(BaseModule):
    def __init__(self, dim, heads=4, dim_head=32):
        super(LinearAttention, self).__init__()
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = torch.nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
        self.to_out = torch.nn.Conv2d(hidden_dim, dim, 1)            

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x)
        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', 
                            heads = self.heads, qkv=3)            
        k = k.softmax(dim=-1)
        context = torch.einsum('bhdn,bhen->bhde', k, v)
        out = torch.einsum('bhde,bhdn->bhen', context, q)
        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', 
                        heads=self.heads, h=h, w=w)
        return self.to_out(out)


class Residual(BaseModule):
    def __init__(self, fn):
        super(Residual, self).__init__()
        self.fn = fn

    def forward(self, x, *args, **kwargs):
        output = self.fn(x, *args, **kwargs) + x
        return output


class SinusoidalPosEmb(BaseModule):
    def __init__(self, dim):
        super(SinusoidalPosEmb, self).__init__()
        self.dim = dim

    def forward(self, x, scale=1000):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb


class GradLogPEstimator2d(BaseModule):
    def __init__(self, dim, dim_mults=(1, 2, 4), groups=8, n_feats=80, pe_scale=1000):
        super(GradLogPEstimator2d, self).__init__()
        self.dim = dim
        self.dim_mults = dim_mults
        self.groups = groups
        self.pe_scale = pe_scale
        
        self.time_pos_emb = SinusoidalPosEmb(dim)
        self.mlp = torch.nn.Sequential(torch.nn.Linear(dim, dim * 4), Mish(),
                                       torch.nn.Linear(dim * 4, dim))

        dims = [2, *map(lambda m: dim * m, dim_mults)]
        in_out = list(zip(dims[:-1], dims[1:]))
        self.downs = torch.nn.ModuleList([])
        self.ups = torch.nn.ModuleList([])
        num_resolutions = len(in_out)

        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (num_resolutions - 1)
            self.downs.append(torch.nn.ModuleList([
                       ResnetBlock(dim_in, dim_out, time_emb_dim=dim),
                       ResnetBlock(dim_out, dim_out, time_emb_dim=dim),
                       Residual(Rezero(LinearAttention(dim_out))),
                       Downsample(dim_out) if not is_last else torch.nn.Identity()]))

        mid_dim = dims[-1]
        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim)
        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim=dim)

        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
            self.ups.append(torch.nn.ModuleList([
                     ResnetBlock(dim_out * 2, dim_in, time_emb_dim=dim),
                     ResnetBlock(dim_in, dim_in, time_emb_dim=dim),
                     Residual(Rezero(LinearAttention(dim_in))),
                     Upsample(dim_in)]))
        self.final_block = Block(dim, dim)
        self.final_conv = torch.nn.Conv2d(dim, 1, 1)

    def forward(self, x, mu, t):
        t = self.time_pos_emb(t, scale=self.pe_scale)
        t = self.mlp(t)

        x = torch.stack([mu, x], 1)

        hiddens = []
        for resnet1, resnet2, attn, downsample in self.downs:
            x = resnet1(x, t)
            x = resnet2(x, t)
            x = attn(x)
            hiddens.append(x)
            x = downsample(x)

        x = self.mid_block1(x, t)
        x = self.mid_attn(x)
        x = self.mid_block2(x, t)

        for resnet1, resnet2, attn, upsample in self.ups:
            x = torch.cat((x, hiddens.pop()), dim=1)
            x = resnet1(x, t)
            x = resnet2(x, t)
            x = attn(x)
            x = upsample(x)

        x = self.final_block(x)
        output = self.final_conv(x)

        return (output).squeeze(1)


class SEBasicBlock(BaseModule):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None, reduction=8):
        super(SEBasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.se = SELayer(planes, reduction)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.relu(out)
        out = self.bn1(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.se(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)
        return out

class SELayer(BaseModule):
    def __init__(self, channel, reduction=8):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
                nn.Linear(channel, channel // reduction),
                nn.ReLU(inplace=True),
                nn.Linear(channel // reduction, channel),
                nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y

class Mish(BaseModule):
    def __init__(self):
        super(Mish, self).__init__()
    def forward(self, x):
        return x * torch.tanh(F.softplus(x))

class MultiHeadAttention(BaseModule):
    ''' Multi-Head Attention module '''
    def __init__(self, n_head, d_model, d_k, d_v, dropout=0., spectral_norm=False):
        super().__init__()
        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k)
        self.w_ks = nn.Linear(d_model, n_head * d_k)
        self.w_vs = nn.Linear(d_model, n_head * d_v)
        
        self.attention = ScaledDotProductAttention(temperature=np.power(d_model, 0.5), dropout=dropout)

        self.layer_norm = nn.LayerNorm(d_model)

        self.fc = nn.Linear(n_head * d_v, d_model)
        self.dropout = nn.Dropout(dropout)

        if spectral_norm:
            self.w_qs = nn.utils.spectral_norm(self.w_qs)
            self.w_ks = nn.utils.spectral_norm(self.w_ks)
            self.w_vs = nn.utils.spectral_norm(self.w_vs)
            self.fc = nn.utils.spectral_norm(self.fc)

    def forward(self, x):
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_x, _ = x.size() 
        
        residual = x
        
        q = self.w_qs(x).view(sz_b, len_x, n_head, d_k)
        k = self.w_ks(x).view(sz_b, len_x, n_head, d_k)
        v = self.w_vs(x).view(sz_b, len_x, n_head, d_v)

        q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_x, d_k)  # (n*b) x lv x dv
        k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_x, d_k)  # (n*b) x lq x dk
        v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_x, d_v)  # (n*b) x lq x dk

        output = self.attention(q, k, v)

        output = output.view(n_head, sz_b, len_x, d_v)
        output = output.permute(1, 2, 0, 3).contiguous().view(sz_b, len_x, -1)  # b x lq x (n*dv)

        output = self.dropout(self.fc(output))
        output = self.layer_norm(output + residual)
        return output


class ScaledDotProductAttention(BaseModule):
    ''' Scaled Dot-Product Attention '''
    def __init__(self, temperature, dropout):
        super().__init__()
        self.temperature = temperature
        self.softmax = nn.Softmax(dim=2)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v):

        attn = torch.bmm(q, k.transpose(1, 2))
        attn = attn / self.temperature

        attn = self.softmax(attn)
        p_attn = self.dropout(attn)

        output = torch.bmm(p_attn, v)
        return output

class PositionwiseFeedForward(BaseModule):
    ''' A two-feed-forward-layer module '''
    def __init__(self, d_in, d_hid, fft_conv1d_kernel_size, dropout=0.1):
        super().__init__()
        self.w_1 = torch.nn.Conv1d(d_in, d_hid, kernel_size=fft_conv1d_kernel_size[0], padding=int((fft_conv1d_kernel_size[0] - 1) / 2))
        self.w_2 = torch.nn.Conv1d(d_hid, d_in, kernel_size=fft_conv1d_kernel_size[1], padding=int((fft_conv1d_kernel_size[1] - 1) / 2))

        self.layer_norm = nn.LayerNorm(d_in)

        self.mish = Mish()
        self.dropout = nn.Dropout(dropout)

    def forward(self, input):
        residual = input

        output = input.transpose(1, 2)
        output = self.w_2(self.dropout(self.mish(self.w_1(output))))
        output = output.transpose(1, 2)

        output = self.dropout(output)
        output = self.layer_norm(output + residual)

        return output
code/diff_sv/models/modules.py
wngh1187/Diff-SV
import torch
import torch.nn as nn
import numpy as np
from models.modules import Mish, MultiHeadAttention, PositionwiseFeedForward
from models.base import BaseModule

class Transformer(BaseModule):
    """ Transformer """
    def __init__(self, n_layers, d_model, n_head=2, d_inner=128):
        super(Transformer, self).__init__()
        self.n_layers = n_layers
        self.d_model = d_model
        self.n_head = n_head
        self.d_k = self.d_model // self.n_head
        self.d_v = self.d_model // self.n_head
        self.d_inner = d_inner
        self.fft_conv1d_kernel_size = [3, 3]
        self.dropout = 0.1
        
        self.spectral = nn.Sequential(
            nn.Linear(self.d_model, self.d_model),
            Mish(),
            nn.Dropout(0.1),
            nn.Linear(self.d_model, self.d_model),
            Mish(),
            nn.Dropout(0.1)
        )

        self.layer_stack = nn.ModuleList([FFTBlock(
            self.d_model, self.d_inner, self.n_head, self.d_k, self.d_v, 
            self.fft_conv1d_kernel_size, self.dropout) for _ in range(self.n_layers)])

    def forward(self, x):
        x = self.spectral(x.transpose(1, 2))

        for enc_layer in self.layer_stack:
            x = enc_layer(x)
            
        return x.transpose(1, 2)

class FFTBlock(BaseModule):
    ''' FFT Block '''
    def __init__(self, d_model, d_inner, n_head, d_k, d_v, fft_conv1d_kernel_size, dropout):
        super(FFTBlock, self).__init__()
        self.slf_attn = MultiHeadAttention(
            n_head, d_model, d_k, d_v, dropout=dropout, spectral_norm=False)

        self.pos_ffn = PositionwiseFeedForward(
            d_model, d_inner, fft_conv1d_kernel_size, dropout=dropout)

    def forward(self, input):
        # multi-head self attn
        output = self.slf_attn(input)

        # position wise FF
        output = self.pos_ffn(output)

        return output



code/diff_sv/models/transformer.py
wngh1187/Diff-SV
import math
from torch.optim.lr_scheduler import _LRScheduler

class CosineAnnealingWarmUpRestarts(_LRScheduler):
    def __init__(self, optimizer, T_0, T_mult=1, eta_max=0.1, T_up=0, gamma=1., last_epoch=-1):
        if T_0 <= 0 or not isinstance(T_0, int):
            raise ValueError("Expected positive integer T_0, but got {}".format(T_0))
        #if T_mult < 1 or not isinstance(T_mult, int):
        if T_mult < 1:
            raise ValueError("Expected integer T_mult >= 1, but got {}".format(T_mult))
        if T_up < 0 or not isinstance(T_up, int):
            raise ValueError("Expected positive integer T_up, but got {}".format(T_up))
        self.T_0 = T_0
        self.T_mult = T_mult
        self.base_eta_max = eta_max
        self.eta_max = eta_max
        self.T_up = T_up
        self.T_i = T_0
        self.gamma = gamma
        self.cycle = 0
        self.T_cur = last_epoch
        super(CosineAnnealingWarmUpRestarts, self).__init__(optimizer, last_epoch)
    
    def get_lr(self):
        if self.T_cur == -1:
            return self.base_lrs
        elif self.T_cur < self.T_up:
            return [(self.eta_max - base_lr)*self.T_cur / self.T_up + base_lr for base_lr in self.base_lrs]
        else:
            return [base_lr + (self.eta_max - base_lr) * (1 + math.cos(math.pi * (self.T_cur-self.T_up) / (self.T_i - self.T_up))) / 2
                    for base_lr in self.base_lrs]

    def step(self, epoch=None):
        if epoch is None:
            epoch = self.last_epoch + 1
            self.T_cur = self.T_cur + 1
            if self.T_cur >= self.T_i:
                self.cycle += 1
                self.T_cur = self.T_cur - self.T_i
                self.T_i = int((self.T_i - self.T_up) * self.T_mult + self.T_up)
        else:
            if epoch >= self.T_0:
                if self.T_mult == 1:
                    self.T_cur = epoch % self.T_0
                    self.cycle = epoch // self.T_0
                else:
                    n = int(math.log((epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult))
                    self.cycle = n
                    self.T_cur = int(epoch - self.T_0 * (self.T_mult ** n - 1) / (self.T_mult - 1))
                    self.T_i = int(self.T_0 * self.T_mult ** (n))
            else:
                self.T_i = self.T_0
                self.T_cur = epoch
                
        self.eta_max = self.base_eta_max * (self.gamma**self.cycle)
        self.last_epoch = math.floor(epoch)
        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
            param_group['lr'] = lr
code/diff_sv/optimizer/CosineAnnealingWarmUpRestarts.py
wngh1187/Diff-SV
import torch
import torchaudio

class LogMelspectrogram():
    """Extract Log-Melspectrogram from raw waveform using torchaudio.
    Note that this module automatically synchronizes device with input tensor.
    """
    def __init__(self, winlen, winstep, nfft, samplerate, nfilts, premphasis, winfunc):
        super(LogMelspectrogram, self).__init__()
        
        self.device = 'cpu'
        self.premphasis = premphasis
        self.mel_spec = torchaudio.transforms.MelSpectrogram(
            n_mels=nfilts,
            n_fft=nfft,
            win_length=winlen,
            hop_length=winstep,
            sample_rate=samplerate,
            window_fn=winfunc
        )

    def __call__(self, x):
        # synchronize device
        if self.device != x.device:
            self.device = x.device
            self.mel_spec.to(x.device)
            
        # pre-emphasis
        if self.premphasis is not None:
            x = x[:, 1:] - self.premphasis * x[:, 0:-1]
        
        
        # melspectrogram
        x = x.type(torch.float32)
        x = self.mel_spec(x)

        # log-mel
        x = x + 1e-6
        x = torch.log(x)

        return x
code/diff_sv/speech_features/log_melspectrogram.py
wngh1187/Diff-SV
from tqdm import tqdm

import torch
import torch.nn.functional as F
import numpy as np
import torch.distributed as dist

from utils.ddp_util import all_gather
import utils.metric as metric
from torch.autograd import Variable

def load_dict(args, model, dic):
    if len(args['gpu_ids']) <=1 and list(dic.keys())[0][:6] == 'module':
        new_dic = {}
        for key, value in zip(dic.keys(), dic.values()):
            new_dic[key[7:]] = value
        model.load_state_dict(new_dic)
    else: model.load_state_dict(dic)
    return model
    
class ModelTrainer:
    def __init__(self, args):
        args = None
        vox1 = None
        voices = None
        model = None
        logger = None
        criterion = None
        optimizer = None
        lr_scheduler = None
        train_loader = None
        enrollment_set = None
        enrollment_loader = None
        dev_loader_voices = None
        eval_loader_voices = None

    def run(self):
        self.best_eer = 1000
        self.update_best_model_dic = [0]
        self.best_model_dic = None
        self.best_epoch = None
        self.idx_for_log = len(self.train_loader) // self.args['number_iteration_for_log']
        self.log_step = 0
        self.global_step = 0
        self._loss = 0.
        self._loss_clf = 0.
        self._loss_prior = 0.
        self._loss_diff = 0.

        for epoch in range(self.args['epoch']):
            self.train_set_sampler.set_epoch(epoch)
            self.train(epoch)
            self.test(epoch)

        ### Activate to evaluate the model using pre-trained weights ###  
        # path_bestmodel =  '/workspace/Diff-SV/code/weight/diff-sv.pt'
        # self.best_model_dic = torch.load(path_bestmodel, map_location=self.args['device'])
        # self.model = load_dict(self.args, self.model, self.best_model_dic['model'])
        
        self.model.load_state_dict(self.best_model_dic)
        self.test_noisy()    

    def train(self, epoch):
        self.model.train()
        idx_ct_start = len(self.train_loader)*(int(epoch))
        with tqdm(total = len(self.train_loader), ncols = 200) as pbar:
            for idx, (m_batch, m_batch_referance, m_label, m_aug) in enumerate(self.train_loader):
                loss = 0
                self.optimizer.zero_grad()
                
                m_label = m_label.to(self.args['device'])
                m_aug = m_aug.to(self.args['device'])
                m_batch = m_batch.to(torch.float32).to(self.args['device'], non_blocking=True)
                m_batch_referance = m_batch_referance.to(torch.float32).to(self.args['device'], non_blocking=True)

                code, loss_prior, loss_diff = self.model(m_batch, m_batch_referance, m_aug, epoch, is_train=True)
                
                description = '%s epoch: %d '%(self.args['name'], epoch)
                
                loss_clf = self.criterion['classification_loss'](code, m_label)
                loss += loss_clf
                self._loss_clf += loss_clf.cpu().detach() 
                description += 'loss_clf:%.3f '%(loss_clf)
                
                loss += loss_prior
                self._loss_prior += loss_prior.cpu().detach() 
                description += 'loss_prior:%.3f '%(loss_prior)

                loss += loss_diff
                self._loss_diff += loss_diff.cpu().detach()
                description += 'loss_diff:%.3f '%(loss_diff)

                self._loss += loss.cpu().detach()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 2.0)
                loss.backward()
                self.optimizer.step()
    
                description += 'TOT: %.4f'%(loss)
                pbar.set_description(description)
                pbar.update(1)
                self.global_step +=1

                ###  if the current epoch is match to the logging condition, log ###
                if self.global_step % self.idx_for_log == 0:
                    self._loss /= self.idx_for_log
                    self._loss_clf /= self.idx_for_log
                    self._loss_diff /= self.idx_for_log
                    self._loss_prior /= self.idx_for_log

                    for p_group in self.optimizer.param_groups:
                        lr = p_group['lr']
                        break

                    if self.args['flag_parent']:
                        self.logger.log_metric('loss', self._loss, log_step = self.log_step)
                        self.logger.log_metric('loss_clf', self._loss_clf, log_step = self.log_step)
                        self.logger.log_metric('loss_diff', self._loss_diff, log_step = self.log_step)
                        self.logger.log_metric('loss_prior', self._loss_prior, log_step = self.log_step)
                        self.logger.log_metric('lr', lr, log_step = self.log_step)

                        self._loss = 0.
                        self._loss_clf = 0.
                        self._loss_diff = 0.
                        self._loss_prior = 0.
                        self.log_step += 1

                self.lr_scheduler.step()        


    def test(self, epoch):
        ### evaluate clean evaluation ###
        self.enrollment_set.Key = 'clean'
        self.embeddings = self._enrollment(self.enrollment_loader)
        if self.args['flag_parent']:
            self.cur_eer, min_dcf = self._calculate_eer(self.vox1.test_trials)
            self.logger.log_metric('EER_clean', self.cur_eer, epoch_step=epoch)
            self.logger.log_metric('Min_DCF_clean', min_dcf, epoch_step=epoch)
    
            if self.cur_eer < self.best_eer:
                self.best_eer = self.cur_eer
                self.logger.log_metric('BestEER_clean', self.best_eer, epoch_step=epoch)
                check_point = {
                    'epoch': epoch,
                    'model': self.model.state_dict(),
                    'optimizer': self.optimizer.state_dict(),
				    'lr_scheduler': self.lr_scheduler.state_dict(),
                }
                self.logger.save_model('Checkpoint_{}_{}'.format(epoch, self.best_eer), check_point)
                self.update_best_model_dic = [1]
        if len(self.args['gpu_ids']) >1: self._synchronize()

        self.update_best_model_dic = all_gather(self.update_best_model_dic)
        if sum(self.update_best_model_dic):
            self.best_model_dic = self.model.state_dict()
            self.best_epoch = epoch
            self.update_best_model_dic = [0]
    
    def test_noisy(self):
        ### 1. in-domain speech evaluation with in-domain noise ###
        self.test_noise_set(self.best_epoch, 'noise_0')
        self.test_noise_set(self.best_epoch, 'noise_5')
        self.test_noise_set(self.best_epoch, 'noise_10')
        self.test_noise_set(self.best_epoch, 'noise_15')
        self.test_noise_set(self.best_epoch, 'noise_20')
        
        self.test_noise_set(self.best_epoch, 'speech_0')
        self.test_noise_set(self.best_epoch, 'speech_5')
        self.test_noise_set(self.best_epoch, 'speech_10')
        self.test_noise_set(self.best_epoch, 'speech_15')
        self.test_noise_set(self.best_epoch, 'speech_20')
        
        self.test_noise_set(self.best_epoch, 'music_0')
        self.test_noise_set(self.best_epoch, 'music_5')
        self.test_noise_set(self.best_epoch, 'music_10')
        self.test_noise_set(self.best_epoch, 'music_15')
        self.test_noise_set(self.best_epoch, 'music_20')

        ### 2. in-domain speech evaluation with out-of-domain noise ###
        self.test_noise_set(self.best_epoch, 'nonspeech_0')
        self.test_noise_set(self.best_epoch, 'nonspeech_5')
        self.test_noise_set(self.best_epoch, 'nonspeech_10')
        self.test_noise_set(self.best_epoch, 'nonspeech_15')
        self.test_noise_set(self.best_epoch, 'nonspeech_20')

        ### 3. out-of-domain speech evaluation with out-of-domain noise ###
        self.test_voices(self.best_epoch, 'voices_dev', self.dev_loader_voices, self.voices.dev_trials)
        self.test_voices(self.best_epoch, 'voices_eval', self.eval_loader_voices, self.voices.eval_trials)

    def test_noise_set(self, epoch, key):
        self.enrollment_set.Key = key
        self.embeddings = self._enrollment(self.enrollment_loader)
        if self.args['flag_parent']:
            eer, min_dcf = self._calculate_eer(self.vox1.test_trials)
            self.logger.log_metric(f'EER_{key}', eer, epoch_step=epoch)
            self.logger.log_metric(f'Min_DCF_{key}', min_dcf, epoch_step=epoch)
        if len(self.args['gpu_ids']) >1: self._synchronize()

    def test_voices(self, epoch, key, loader, trial):
        self.embeddings = self._enrollment(loader)
        if self.args['flag_parent']:
            eer, min_dcf = self._calculate_eer(trial)
            self.logger.log_metric(f'EER_{key}', eer, epoch_step=epoch)
            self.logger.log_metric(f'Min_DCF_{key}', min_dcf, epoch_step=epoch)
        if len(self.args['gpu_ids']) >1: self._synchronize()

    def _enrollment(self, enrollment_loader):
        ### Return embedding dictionary ###
        self.model.eval()

        keys = []
        embeddings = []
        
        with tqdm(total=len(enrollment_loader), ncols=90) as pbar, torch.set_grad_enabled(False):
            if len(self.args['gpu_ids']) >1:
                with self.model.no_sync():
                    for x_seg, key in enrollment_loader:

                        x_seg = x_seg.to(torch.float32).to(self.args['device'], non_blocking=True).view(-1, x_seg.size(-1)) 
                        x_seg = self.model(x_seg, is_train=False).to('cpu').view(-1, self.args['num_seg'], self.args['code_dim']) 

                        keys.extend(key)
                        embeddings.extend(x_seg)

                        if self.args['flag_parent']:
                            pbar.update(1)
                            
                keys = all_gather(keys)
                embeddings = all_gather(embeddings)
            else:
                for x_seg, key in enrollment_loader:

                    x_seg = x_seg.to(torch.float32).to(self.args['device'], non_blocking=True).view(-1, x_seg.size(-1)) 
                    x_seg = self.model(x_seg, is_train=False).to('cpu').view(-1, self.args['num_seg'], self.args['code_dim']) 

                    keys.extend(key)
                    embeddings.extend(x_seg)
                    pbar.update(1)

        seg_dict = {}
        for i in range(len(keys)):
            seg_dict[keys[i]] = embeddings[i]

        return seg_dict
    
    def _calculate_eer(self, test_trials):
        labels = []
        cos_sims = [[], []]

        for item in tqdm(test_trials, desc='test', ncols=self.args['tqdm_ncols']):
            cos_sims[0].append(self.embeddings[item.key1])
            cos_sims[1].append(self.embeddings[item.key2])
            labels.append(int(item.label))
        
        batch = len(labels)
        num_seg = self.args['num_seg']
        buffer1 = torch.stack(cos_sims[0], dim=0).view(batch, num_seg, -1)
        buffer2 = torch.stack(cos_sims[1], dim=0).view(batch, num_seg, -1)
        buffer1 = buffer1.repeat(1, num_seg, 1).view(batch * num_seg * num_seg, -1)
        buffer2 = buffer2.repeat(1, 1, num_seg).view(batch * num_seg * num_seg, -1)
        cos_sims = F.cosine_similarity(buffer1, buffer2)
        cos_sims = cos_sims.view(batch, num_seg * num_seg)
        cos_sims = cos_sims.mean(dim=1)

        eer = metric.calculate_EER(
            scores=cos_sims, labels=labels
        )
        min_dcf = metric.calculate_MinDCF(
            scores=cos_sims, labels=labels
        )
        return eer, min_dcf

    def _synchronize(self):
        torch.cuda.empty_cache()
        dist.barrier()
code/diff_sv/trainers/train.py
wngh1187/Diff-SV
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
This file contains primitives for multi-gpu communication.
This is useful when doing distributed training.
"""

import functools
import logging
import numpy as np
import pickle
import torch
import torch.distributed as dist

_LOCAL_PROCESS_GROUP = None
"""
A torch process group which only includes processes that on the same machine as the current process.
This variable is set when processes are spawned by `launch()` in "engine/launch.py".
"""

class GatherLayer(torch.autograd.Function):
    '''Gather tensors from all process, supporting backward propagation.
    '''

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        output = [torch.zeros_like(input) \
            for _ in range(dist.get_world_size())]
        dist.all_gather(output, input)
        return tuple(output)

    @staticmethod
    def backward(ctx, *grads):
        input, = ctx.saved_tensors
        grad_out = torch.zeros_like(input)
        grad_out[:] = grads[dist.get_rank()]
        return grad_out


def get_world_size() -> int:
	if not dist.is_available():
		return 1
	if not dist.is_initialized():
		return 1
	return dist.get_world_size()


def get_rank() -> int:
	if not dist.is_available():
		return 0
	if not dist.is_initialized():
		return 0
	return dist.get_rank()


def get_local_rank() -> int:
	"""
	Returns:
		The rank of the current process within the local (per-machine) process group.
	"""
	if not dist.is_available():
		return 0
	if not dist.is_initialized():
		return 0
	assert _LOCAL_PROCESS_GROUP is not None
	return dist.get_rank(group=_LOCAL_PROCESS_GROUP)


def get_local_size() -> int:
	"""
	Returns:
		The size of the per-machine process group,
		i.e. the number of processes per machine.
	"""
	if not dist.is_available():
		return 1
	if not dist.is_initialized():
		return 1
	return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)


def is_main_process() -> bool:
	return get_rank() == 0


def synchronize():
	"""
	Helper function to synchronize (barrier) among all processes when
	using distributed training
	"""
	if not dist.is_available():
		return
	if not dist.is_initialized():
		return
	world_size = dist.get_world_size()
	if world_size == 1:
		return
	dist.barrier()


@functools.lru_cache()
def _get_global_gloo_group():
	"""
	Return a process group based on gloo backend, containing all the ranks
	The result is cached.
	"""
	if dist.get_backend() == "nccl":
		return dist.new_group(backend="gloo")
	else:
		return dist.group.WORLD


def _serialize_to_tensor(data, group):
	backend = dist.get_backend(group)
	assert backend in ["gloo", "nccl"]
	device = torch.device("cpu" if backend == "gloo" else "cuda")

	buffer = pickle.dumps(data)
	if len(buffer) > 1024 ** 3:
		logger = logging.getLogger(__name__)
		logger.warning(
			"Rank {} trying to all-gather {:.2f} GB of data on device {}".format(
				get_rank(), len(buffer) / (1024 ** 3), device
			)
		)
	storage = torch.ByteStorage.from_buffer(buffer)
	tensor = torch.ByteTensor(storage).to(device=device)
	return tensor


def _pad_to_largest_tensor(tensor, group):
	"""
	Returns:
		list[int]: size of the tensor, on each rank
		Tensor: padded tensor that has the max size
	"""
	world_size = dist.get_world_size(group=group)
	assert (
		world_size >= 1
	), "comm.gather/all_gather must be called from ranks within the given group!"
	local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device)
	size_list = [
		torch.zeros([1], dtype=torch.int64, device=tensor.device) for _ in range(world_size)
	]
	dist.all_gather(size_list, local_size, group=group)
	size_list = [int(size.item()) for size in size_list]

	max_size = max(size_list)

	# we pad the tensor because torch all_gather does not support
	# gathering tensors of different shapes
	if local_size != max_size:
		padding = torch.zeros((max_size - local_size,), dtype=torch.uint8, device=tensor.device)
		tensor = torch.cat((tensor, padding), dim=0)
	return size_list, tensor


def all_gather(data, group=None, append = False):
	"""
	Run all_gather on arbitrary picklable data (not necessarily tensors).
	Args:
		data: any picklable object
		group: a torch process group. By default, will use a group which
			contains all ranks on gloo backend.
	Returns:
		list[data]: list of data gathered from each rank
	"""
	if get_world_size() == 1:
		return [data]
	if group is None:
		group = _get_global_gloo_group()
	if dist.get_world_size(group) == 1:
		return [data]

	tensor = _serialize_to_tensor(data, group)

	size_list, tensor = _pad_to_largest_tensor(tensor, group)
	max_size = max(size_list)

	# receiving Tensor from all ranks
	tensor_list = [
		torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
	]
	dist.all_gather(tensor_list, tensor, group=group)
	data_list = []
	for size, tensor in zip(size_list, tensor_list):
		buffer = tensor.cpu().numpy().tobytes()[:size]
		if append:
			data_list.append(pickle.loads(buffer))
		else:
			data_list.extend(pickle.loads(buffer))
	
	return data_list


def all_gather_tensor(data, gpu, group=None):
	"""
	Run all_gather on arbitrary picklable data (not necessarily tensors).
	Args:
		data: any picklable object
		group: a torch process group. By default, will use a group which
			contains all ranks on gloo backend.
	Returns:
		list[data]: list of data gathered from each rank
	"""
	if get_world_size() == 1:
		return [data]
	if group is None:
		group = _get_global_gloo_group()
	if dist.get_world_size(group) == 1:
		return [data]

	tensor = _serialize_to_tensor(data, group)

	size_list, tensor = _pad_to_largest_tensor(tensor, group)
	max_size = max(size_list)

	# receiving Tensor from all ranks
	tensor_list = [
		torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
	]
	dist.all_gather(tensor_list, tensor, group=group)
	data_list = []
	for size, tensor in zip(size_list, tensor_list):
		buffer = tensor.cpu().numpy().tobytes()[:size]
		#data_list.append(pickle.loads(buffer))
		data_list.extend(pickle.loads(buffer).cuda(gpu, non_blocking = True))
	data_list = torch.stack(data_list)
	return data_list

def gather(data, dst=0, group=None):
	"""
	Run gather on arbitrary picklable data (not necessarily tensors).
	Args:
		data: any picklable object
		dst (int): destination rank
		group: a torch process group. By default, will use a group which
			contains all ranks on gloo backend.
	Returns:
		list[data]: on dst, a list of data gathered from each rank. Otherwise,
			an empty list.
	"""
	if get_world_size() == 1:
		return [data]
	if group is None:
		group = _get_global_gloo_group()
	if dist.get_world_size(group=group) == 1:
		return [data]
	rank = dist.get_rank(group=group)

	tensor = _serialize_to_tensor(data, group)
	size_list, tensor = _pad_to_largest_tensor(tensor, group)

	# receiving Tensor from all ranks
	if rank == dst:
		max_size = max(size_list)
		tensor_list = [
			torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
		]
		dist.gather(tensor, tensor_list, dst=dst, group=group)

		data_list = []
		for size, tensor in zip(size_list, tensor_list):
			buffer = tensor.cpu().numpy().tobytes()[:size]
			data_list.append(pickle.loads(buffer))
		return data_list
	else:
		dist.gather(tensor, [], dst=dst, group=group)
		return []

# def broadcast(data, dst=0, group=None):
# 	if group is None:
# 		group = _get_global_gloo_group()
# 	dist.broadcast(data, dst, group)

def shared_random_seed():
	"""
	Returns:
		int: a random number that is the same across all workers.
			If workers need a shared RNG, they can use this shared seed to
			create one.
	All workers must call this function, otherwise it will deadlock.
	"""
	ints = np.random.randint(2 ** 31)
	all_ints = all_gather(ints)
	return all_ints[0]


def reduce_dict(input_dict, average=True):
	"""
	Reduce the values in the dictionary from all processes so that process with rank
	0 has the reduced results.
	Args:
		input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
		average (bool): whether to do average or sum
	Returns:
		a dict with the same keys as input_dict, after reduction.
	"""
	world_size = get_world_size()
	if world_size < 2:
		return input_dict
	with torch.no_grad():
		names = []
		values = []
		# sort the keys so that they are consistent across processes
		for k in sorted(input_dict.keys()):
			names.append(k)
			values.append(input_dict[k])
		values = torch.stack(values, dim=0)
		dist.reduce(values, dst=0)
		if dist.get_rank() == 0 and average:
			# only main process gets accumulated, so only divide by
			# world_size in this case
			values /= world_size
		reduced_dict = {k: v for k, v in zip(names, values)}
	return reduced_dict

code/diff_sv/utils/ddp_util.py
wngh1187/Diff-SV
import os
import numpy as np
from PIL import Image

import matplotlib.pyplot as plt
from sklearn import metrics
from scipy.optimize import brentq
from scipy.interpolate import interp1d

def calculate_EER(scores, labels):
    if len(scores) != len(labels):
        raise Exception('length between scores and labels is different')
    elif len(scores) == 0:
        raise Exception("There's no elements in scores")
        
    fpr, tpr, _ = metrics.roc_curve(labels, scores, pos_label=1)
    EER = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)

    return EER

def calculate_MinDCF(scores, labels, p_target=0.01, c_miss=10, c_false_alarm=1):
    if len(scores) != len(labels):
        raise Exception('length between scores and labels is different')
    elif len(scores) == 0:
        raise Exception("There's no elements in scores")
    fpr, tpr, _ = metrics.roc_curve(labels, scores, pos_label=1)
    fnr = 1 - tpr

    dcf = c_miss * fnr * p_target + c_false_alarm * fpr * (1 - p_target)
    c_det = np.min(dcf)
    c_def = min(c_miss * p_target, c_false_alarm * (1 - p_target))
    return c_det / c_def

def draw_histogram(scores, labels):
    positive = []
    negative = []
    for i in range(len(labels)):
        if int(labels[i]) == 1:
            positive.append(scores[i])
        else:
            negative.append(scores[i])
    
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    plt.ylim([0, 500])
    ax.hist(positive, label="target", bins=200, color="blue", alpha=0.5)
    ax.hist(negative, label="nontarget", bins=200, color="red", alpha=0.5)
    ax.xaxis.set_ticks_position("bottom")
    ax.yaxis.set_ticks_position("left")
    plt.ylabel("#trial")
    plt.legend(loc="best")

    if os.path.exists('histogram.png'):
        os.remove('histogram.png')    

    plt.savefig("histogram.png", dpi=400, bbox_inches="tight")
    plt.close(fig)
    img = Image.open("histogram.png")
    
    return img

def draw_heatmap(title, x_axis, y_axis, featmap):
    plt.matshow(featmap, cmap=plt.get_cmap('GnBu'), aspect='auto')
    plt.colorbar(shrink=0.8, aspect=10)
    plt.clim(np.min(featmap), np.max(featmap))
    plt.title(title, fontsize=20)
    plt.xlabel(x_axis, fontsize=14)
    plt.ylabel(y_axis, fontsize=14)
    plt.savefig("heatmap.png", dpi=400, bbox_inches="tight")
    plt.close()
    img = Image.open("heatmap.png")
    os.remove('heatmap.png')

    return img
code/diff_sv/utils/metric.py
wngh1187/Diff-SV
import math
import random
import numpy as np

import torch
import torch.nn as nn

__all__=['duplicate', 'subtensor', 'linspace_crop', 'rand_crop']

def duplicate(x, size, dim=-1):
    """duplicate tensor in given dimension 
    until x is larger than size

    params
        x       - tensor to duplicate
        size    - minimum size
        dim     - target dimension
    """
    length = x.size(dim)
    x = torch.cat(tuple(
        x for i in range(math.ceil(size / length))), dim=dim)
    return x 
        
def subtensor(x, start, size, dim=-1):
    """    
    Param
        x       - tensor to crop
        start   - start index
        size    - size of tensor
        dim     - target dimension
    """
    dim_indexes = [i for i in range(len(x.size()))]
    temp = dim_indexes[0]
    dim_indexes[0] = dim_indexes[dim]
    dim_indexes[dim] = temp

    x = x.permute(*dim_indexes) # [..., default_size] -> [defalut_size, ...]
    x = x[start:start + size]
    x = x.permute(*dim_indexes) # [defalut_size, ...] -> [..., default_size]
    
    return x

def linspace_crop(x, num, size, dim=-1):
    """linspace crop for Test Time Augmentation(TTA)

    Param
        x       - tensor to crop
        num     - linspace num
        size    - size of tensor
        dim     - target dimension
    """
    # duplicate
    if x.size(dim) < size:
        x = duplicate(x, size, dim)

    # calculate index
    index = np.linspace(0, x.size(dim) - size, num=num)

    # init buffer
    buffer_size = list(x.size())
    buffer_size[dim] = size
    buffer_size = tuple([len(index)] + buffer_size)
    tensor_buffer = torch.zeros(buffer_size, device='cuda' if x.is_cuda else 'cpu')

    # slice
    for i in range(num):
        tensor_buffer[i] = subtensor(x, int(index[i]), size, dim)

    return tensor_buffer

def rand_crop(x, size, dim=-1):
    """random crop for data augmentation

    Param
        x       - tensor to crop
        size    - size of tensor
        dim     - target dimension
    """
    if x.size(dim) < size:
        x = duplicate(x, size, dim)
    index = random.randint(0, x.size(dim) - size)
    return subtensor(x, index, size, dim)

def init_weights(m):
	if isinstance(m, nn.Linear):
		nn.init.xavier_uniform(m.weight)
		m.bias.data.fill_(0.0001)
	elif isinstance(m, nn.BatchNorm1d) or isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.InstanceNorm1d):
		pass
	else:
		if hasattr(m, 'weight'):
			nn.init.kaiming_normal_(m.weight, a=0.01)
		else:
			pass
code/diff_sv/utils/util.py
wngh1187/Diff-SV