Self-supervised speaker recognition with LGL代码分析

1.model.py

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

import numpy, tqdm, sys, time, soundfile

from loss import *
from encoder import *
from tools import *

class model(nn.Module):
    def __init__(self, lr, lr_decay, **kwargs):
        super(model, self).__init__()
        self.Network = ECAPA_TDNN().cuda() # Speaker encoder
        self.Loss = LossFunction().cuda() # Contrastive loss
        self.AATNet  = AATNet().cuda() # AAT, which is used to improve the performace
        self.Reverse = Reverse().cuda() # AAT
        self.OptimNet = torch.optim.Adam(list(self.Network.parameters()) + list(self.Loss.parameters()), lr = lr)
        self.OptimAAT = torch.optim.Adam(self.AATNet.parameters(), lr = lr)
        self.Scheduler = torch.optim.lr_scheduler.StepLR(self.OptimNet, step_size = 5, gamma=lr_decay)
        print("Model para number = %.2f"%(sum(param.numel() for param in self.Network.parameters()) / 1024 / 1024))

    def train_network(self, loader, epoch):
        # Contrastive learning with AAT, for more details about AAT, please check here: https://github.com/joonson/voxceleb_unsupervised
        self.train()
        self.Scheduler.step(epoch - 1) # Update the learning rate
        loss, top1 = 0, 0
        lr = self.OptimNet.param_groups[0]['lr'] # Read the current learning rate
        criterion   = torch.nn.CrossEntropyLoss() # Use for AAT
        AAT_labels = torch.LongTensor([1]*loader.batch_size+[0]*loader.batch_size).cuda() # AAT labels
        tstart = time.time() # Used to monitor the training speed
        for counter, data in enumerate(loader, start = 1):     
            data = data.transpose(0,1)
            feat = []
            for inp in data:
                feat.append(self.Network.forward(torch.FloatTensor(inp).cuda())) # Feed the segments to get the speaker embeddings
            feat = torch.stack(feat,dim=1).squeeze()
            self.zero_grad()
            # Train discriminator
            out_a, out_s, out_p = feat[:,0,:].detach(), feat[:,1,:].detach(), feat[:,2,:].detach()
            in_AAT = torch.cat((torch.cat((out_a,out_s),1),torch.cat((out_a,out_p),1)),0)
            out_AAT = self.AATNet(in_AAT)
            dloss  = criterion(out_AAT, AAT_labels)
            dloss.backward()
            self.OptimAAT.step()
            # Train model
            self.zero_grad()
            in_AAT = torch.cat((torch.cat((feat[:,0,:],feat[:,1,:]),1),torch.cat((feat[:,0,:],feat[:,2,:]),1)),0)
            out_AAT = self.AATNet(self.Reverse(in_AAT))
            closs   = criterion(out_AAT, AAT_labels) # AAT loss
            sloss, prec1 = self.Loss.forward(feat[:,[0,2],:])  # speaker loss              
            nloss = sloss + closs * 3 # Total loss
            loss    += nloss.detach().cpu()
            top1    += prec1 # Training acc
            nloss.backward()
            self.OptimNet.step() 
            time_used = time.time() - tstart # Time for this epoch
            sys.stdout.write("[%2d] Lr: %5f, %.2f%% (est %.1f mins) Loss %f EER/TAcc %2.3f%% \r"%(epoch, lr, 100 * (counter / loader.__len__()), time_used * loader.__len__() / counter / 60, loss/counter, top1/counter))
            sys.stdout.flush()
        sys.stdout.write("\n")
        return loss/counter, top1/counter, lr

    def evaluate_network(self, val_list, val_path, **kwargs):
        self.eval()
        files, feats = [], {}
        for line in open(val_list).read().splitlines():
            data = line.split()
            files.append(data[1])
            files.append(data[2])
        setfiles = list(set(files))
        setfiles.sort()  # Read the list of wav files
        for idx, file in tqdm.tqdm(enumerate(setfiles), total = len(setfiles)):
            audio, _ = soundfile.read(os.path.join(val_path, file))
            feat = torch.FloatTensor(numpy.stack([audio], axis=0)).cuda()
            with torch.no_grad():
                ref_feat = self.Network.forward(feat).detach().cpu()
            feats[file]     = ref_feat # Extract features for each data, get the feature dict
        scores, labels  = [], []
        for line in open(val_list).read().splitlines():
            data = line.split()
            ref_feat = F.normalize(feats[data[1]].cuda(), p=2, dim=1) # feature 1
            com_feat = F.normalize(feats[data[2]].cuda(), p=2, dim=1) # feature 2
            score = numpy.mean(torch.matmul(ref_feat, com_feat.T).detach().cpu().numpy()) # Get the score
            scores.append(score)
            labels.append(int(data[0]))
        EER = tuneThresholdfromScore(scores, labels, [1, 0.1])[1]
        fnrs, fprs, thresholds = ComputeErrorRates(scores, labels)
        minDCF, _ = ComputeMinDcf(fnrs, fprs, thresholds, 0.05, 1, 1)
        return EER, minDCF

    def save_network(self, path): # Save the model
        torch.save(self.state_dict(), path)

    def load_network(self, path): # Load the parameters of the pretrain model
        self_state = self.state_dict()
        loaded_state = torch.load(path)
        print("Model %s loaded!"%(path))
        for name, param in loaded_state.items():
            origname = name
            if name not in self_state:
                name = name.replace("module.", "")
                if name not in self_state:
                    print("%s is not in the model."%origname)
                    continue
            if self_state[name].size() != loaded_state[origname].size():
                print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, self_state[name].size(), loaded_state[origname].size()))
                continue
            self_state[name].copy_(param)

 1.train_network

在train_network中，这段代码针对输入数据进行操作，首先使用 enumerate 函数给每个 batch添加上 counter 序号。然后将数据 data 的第一维和第二维交换（按照通道数和批次大小的顺序），这是因为反卷积计算时需要交换输入数据的通道维度和批次维度。

在以下 for 循环中，将每个序列输入到网络结构 self.Network.forward 中，返回输出的 speaker embeddings。最后将 speaker embeddings 合并成 Tensor。

在此之后，调用 self.zero_grad()，将网络参数的梯度初始化为0，以便在下一批次的反向传播时使用。这是在使用 PyTorch 训练神经网络时的一个必要步骤。

从 speaker embeddings tensor feat 中提取出每个 batch 的第一个、第二个和第三个（下标为0、1、2）元素，分别赋给 out_a、out_s 和 out_p，并分别将它们的梯度显式地断开(detach)。这通常是因为梯度在此时不需要计算。

接下来，将 out_a、out_s 和 out_p 拼接成一个 2 * K 维的输入张量，其中 K 表示 speaker embedding 的维数。这个输入张量会被送入 AATNet 中进行处理，其中 AATNet是一个神经网络模型，它的输入是由 2 个 speaker embeddings 拼接起来的向量，输出是一个单一的值。

使用损失函数 criterion 计算预测值 out_AAT 与标签 AAT_labels 之间的差异，并将误差反向传播来计算梯度。最后，用 OptimAAT（一个优化器）更新 AATNet 中的参数，以最小化损失。

总之，这段代码的目的是在神经网络中计算嵌入向量的距离，并通过调整模型参数来使它们更好地符合预期的距离(在此处由AAT_labels表示)。这是跟训练有关的代码。

紧接着

这段代码是模型的训练步骤，其中：

首先，通过在模型上调用 self.zero_grad() 将梯度清零。

然后，将 speaker embeddings 拼接起来（与上文类似），送入 AATNet 中进行处理，返回预测结果 out_AAT。在此处用损失函数 criterion 计算 AAT loss (closs)。

接下来，将 speaker embeddings 提取出来，送入 Loss神经网络模型中进行处理，返回预测结果和精准度。精准度指的是模型输出与实际结果相符合的比例（即输出最大值的下标与标签相等的次数占总次数的比例）。

将上述两个误差加权相加，更新整个网络的权重。最后，计算并打印出本次迭代的时间、学习率(lr)、迭代数(counter)，每个epoch中Loss的平均值、训练精度、估计剩余时间等关键信息。

2.evaluate_network

这段代码用于测试一个预训练模型的性能，它的输入参数包括验证数据集列表(val_list)和音频文件夹的路径(val_path)等。

首先，将网络设置为评估模式，即将 Dropout 和 BatchNorm等 regularization 层设置为测试模式。

接下来，打开 "val_list" 中列出的每个音频文件，提取 嵌入embeddings，并将它们保存在一个字典 "feats" 中。

然后，对于每一对验证数据（一个正样本和一个负样本），从字典 "feats" 中提取对应的两个嵌入 embedding，并通过点积计算它们的相似度score。这个 score 在算法的性能评估中很重要。

最后，使用 scores 和 labels 来计算 Equal Error Rate (EER)和 Minimum Detection Cost Function (minDCF)，并将它们作为结果返回。其中，EER 以及 tuneThresholdFromScore 和 ComputeErrorRates 函数都是用来评估系统性能的工具。

首先从一个文本文件中读取音频文件列表，并生成一个按字母顺序排序的文件列表。它接受一个参数“val_list”，该参数是一个文本文件路径，其中包含每个音频文件的标签和文件位置信息，格式如下：

1 /path/to/reference.wav /path/to/test.wav

        该代码将打开文件、读取其中的每一行，然后将一对音频文件（参考和测试）所在的文件名添加到一个名为“files”的列表中。

接下来，通过将“files”列表转换为一个“set”对象，将文件列表中的重复文件名去除。随后，将这个“set”集合转为列表，而后使用“sort()”函数来对列表中的元素进行字母顺序排序，这一步目的是按字母顺序对所有音频文件进行排序以便进一步处理。

最后，该代码返回一个名为“setfiles”的列表，其中包含了去重、排序后的所有音频文件名。这些文件名作为数据集的索引被用于之前的函数——evaluate_network()中。

        使用上面获取的排好序的音频文件列表“setfiles”，循环遍历每个文件夹，读取当前文件夹中的音频文件，以便进行端到端语音识别。

在循环中，首先会从音频文件中读取和存储音频信号。然后，通过使用预训练模型对读入的音频文件进行前向传递，获得音频嵌入。嵌入被附加到“feats”字典中，其中键是音频文件名。

具体来说，代码使用 SoundFile 库从文件系统中读取音频文件，返回音频信号及其采样率。个人理解使用stack函数，会将一维的音频信号数据沿着axis=0的方向堆叠成一个二维矩阵，以便进行处理。为了使用GPU加速，将生成的嵌入向量转换为PyTorch张量，便于在GPU上进行附加操作。

然后，调用预先训练的神经网络（accessible as“self.Network”）的forwards()方法（self.Network.forward(feat))对输入数据进行前向传递，并得到对应的音频嵌入。这个音频嵌入ref_feat是一个PyTorch张量，并使用detach()和cpu()分别从GPU中分离和转移到CPU上。最后，将这些嵌入向量存储在字典“feats”中，以便后续的计算。循环结束后，“feats”字典包含每个音频文件名及其对应的特征嵌入向量，这些特征向量将用于后续的相似度计算。

2.dataloader.py

import torch, numpy, random, os, math, glob, soundfile
from torch.utils.data import Dataset, DataLoader
from scipy import signal

class train_loader(Dataset):
    def __init__(self, max_frames, train_list, train_path, musan_path, **kwargs):
        self.max_frames = max_frames
        self.data_list = []
        self.noisetypes = ['noise','speech','music'] # Type of noise
        self.noisesnr = {'noise':[0,15],'speech':[13,20],'music':[5,15]} # The range of SNR
        self.noiselist = {} 
        augment_files   = glob.glob(os.path.join(musan_path,'*/*/*.wav')) # All noise files in list
        for file in augment_files:
            if not file.split('/')[-3] in self.noiselist:
                self.noiselist[file.split('/')[-3]] = []
            self.noiselist[file.split('/')[-3]].append(file) # All noise files in dic
        self.rir_files = numpy.load('/student/temp/jyt522/Loss-Gated-Learning-main-vox/Stage1/rir.npy') # Load the rir file
        for line in open(train_list).read().splitlines():
            filename = os.path.join(train_path, line.split()[1])
            self.data_list.append(filename) # Load the training data list
                
    def __getitem__(self, index):
        audio = loadWAVSplit(self.data_list[index], self.max_frames).astype(numpy.float) # Load one utterance
        augment_profiles, audio_aug = [], []
        for ii in range(0,2): # Two segments of one utterance
            rir_gains = numpy.random.uniform(-7,3,1)
            rir_filts = random.choice(self.rir_files)
            noisecat    = random.choice(self.noisetypes)
            noisefile   = random.choice(self.noiselist[noisecat].copy()) # Augmentation information for each segment
            snr = [random.uniform(self.noisesnr[noisecat][0],self.noisesnr[noisecat][1])]
            p = random.random()
            if p < 0.25:  # Add rir only
                augment_profiles.append({'rir_filt':rir_filts, 'rir_gain':rir_gains, 'add_noise': None, 'add_snr': None})
            elif p < 0.50: # Add noise only
                augment_profiles.append({'rir_filt':None, 'rir_gain':None, 'add_noise': noisefile, 'add_snr': snr})
            else: # Add both
                augment_profiles.append({'rir_filt':rir_filts, 'rir_gain':rir_gains, 'add_noise': noisefile, 'add_snr': snr})
        audio_aug.append(self.augment_wav(audio[0],augment_profiles[0])) # Segment 0 with augmentation method 0
        audio_aug.append(self.augment_wav(audio[1],augment_profiles[0])) # Segment 1 with augmentation method 0, used for AAT
        audio_aug.append(self.augment_wav(audio[1],augment_profiles[1])) # Segment 1 with augmentation method 1
        audio_aug = numpy.concatenate(audio_aug,axis=0) # Concate and return
        return torch.FloatTensor(audio_aug)

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

    def augment_wav(self,audio,augment):
        if augment['rir_filt'] is not None:
            rir     = numpy.multiply(augment['rir_filt'], pow(10, 0.1 * augment['rir_gain']))    
            audio   = signal.convolve(audio, rir, mode='full')[:len(audio)]
        if augment['add_noise'] is not None:
            noiseaudio  = loadWAV(augment['add_noise'], self.max_frames).astype(numpy.float)
            noise_db = 10 * numpy.log10(numpy.mean(noiseaudio[0] ** 2)+1e-4) 
            clean_db = 10 * numpy.log10(numpy.mean(audio ** 2)+1e-4) 
            noise = numpy.sqrt(10 ** ((clean_db - noise_db - augment['add_snr']) / 10)) * noiseaudio
            audio = audio + noise
        else:
            audio = numpy.expand_dims(audio, 0)
        return audio

def loadWAV(filename, max_frames):
    max_audio = max_frames * 160 + 240 # 240 is for padding, for 15ms since window is 25ms and step is 10ms.
    audio, _ = soundfile.read(filename)
    audiosize = audio.shape[0]
    if audiosize <= max_audio: # Padding if the length is not enough
        shortage    = math.floor( ( max_audio - audiosize + 1 ) / 2 )
        audio       = numpy.pad(audio, (shortage, shortage), 'wrap')
        audiosize   = audio.shape[0]
    startframe = numpy.int64(random.random()*(audiosize-max_audio)) # Randomly select a start frame to extract audio
    feat = numpy.stack([audio[int(startframe):int(startframe)+max_audio]],axis=0)
    return feat

def loadWAVSplit(filename, max_frames): # Load two segments
    max_audio = max_frames * 160 + 240
    audio, _ = soundfile.read(filename)
    audiosize = audio.shape[0]
    if audiosize <= max_audio:
        shortage    = math.floor( ( max_audio - audiosize) / 2 )
        audio       = numpy.pad(audio, (shortage, shortage), 'wrap')
        audiosize   = audio.shape[0]
    randsize = audiosize - (max_audio*2) # Select two segments
    startframe = random.sample(range(0, randsize), 2)
    startframe.sort()
    startframe[1] += max_audio # Non-overlapped two segments
    startframe = numpy.array(startframe)
    numpy.random.shuffle(startframe)
    feats = []
    for asf in startframe: # Startframe[0] means the 1st segment, Startframe[1] means the 2nd segment
        feats.append(audio[int(asf):int(asf)+max_audio])
    feat = numpy.stack(feats,axis=0)
    return feat

def worker_init_fn(worker_id):
    numpy.random.seed(numpy.random.get_state()[1][0] + worker_id)

def get_loader(args): # Define the data loader
    trainLoader = train_loader(**vars(args))
    trainLoader = torch.utils.data.DataLoader(
        trainLoader,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.n_cpu,
        pin_memory=False,
        drop_last=True,
        worker_init_fn=worker_init_fn,
        prefetch_factor=5,
    )
    return trainLoader

def __getitem__

段代码将一个语音片段分割成若干段较短的语音片段，以便进行神经网络训练和数据增强。具体地说，这个函数是一个数据集类，用于从指定的语音数据集中获取数据条目。其中的 __getitem__(self, index) 函数是 Python 数据集接口的一部分，它会从数据集中获取一个指定数字索引的数据样本。

具体来说，loadWAVSplit() 函数将读取数据片段，将它们分割成固定长度的语音片段，并将它们存储在数组 audio 中。然后，在数据增强阶段，代码将两个语音片段上都应用随机抖动和添加不同的噪声进行数据增强操作。增强操作的方法和参数都是不同的，这样可以增加数据样本的多样性和鲁棒性。

最后，这些处理过的语音片段被连结在一起，并转换成 PyTorch Tensor 对象。这些语音片段就是神经网络训练的输入数据

 

def augment_wav(self,audio,augment):

这个函数是用于对输入的语音片段应用数据增强算法的辅助函数。具体来说，当 augment 非空时，函数将尝试实现两种数据增强方法：添加房间脉冲响应（RIR）和添加噪声。增强方法的相关设置被保存在 augment 字典中。具体来说，函数将输入 audio 加上选择的 RIR，并将结果存储在 audio 变量中；然后，如果选择的噪声存在，则按指定信噪比计算噪声强度，并将其叠加到输入的语音片段中。最后，函数将增强过的语音片段返回给调用它的函数。

应用增强方法的技术细节稍微复杂一些，因此需要注意以下几点：

• 增强函数首先检查 augment 字典以查看是否选择了 RIR 或添加噪声，如果没有对应选项，则不进行数据增强处理；
• 如果选择了 RIR，增强函数将加载选择的 RIR 文件，并将 RIR 加权应用到输入语音片段中；
• 如果选择了添加噪声，则增强函数将加载选择的噪声文件，并将其信噪比计算与输入语音片段适当混合，以使所得语音片段的信噪比符合指定要求；
• 最后，增强函数返回增强后的语音片段。

 

def loadWAV(filename, max_frames):

这个函数是一个辅助函数，用于从给定文件中读取和缩放语音片段。具体来说，该函数首先将给定文件读入内存中，并将语音片段缩放到规定的长度（ max_audio）。如果读取的语音片段长度不足 max_audio，函数将在两端进行填充，并将其存储在数组 audio 中。

然后，函数将从 audio 中随机选择一个起始位置（即“语音抖动”）。从该位置开始，函数将使用一个长度为 max_audio 的滑动窗口在语音片段中滑动，并最终选择最具代表性的语音片段。这个语音片段将被存储在大小为 [1, max_audio] 的 NumPy 数组 feat 中，并作为函数的返回值。

需要注意的是，该函数与前面提到的 loadWAVSplit() 函数略有不同。前者是用于数据增强前的初始加载和缩放，而后者则是用于从原始语音文件中生成多个数据增强片段。

 

def loadWAVSplit(filename, max_frames): 

这个函数是一个辅助函数，用于从指定文件中分割两个语音片段。具体来说，函数首先读取给定文件，将语音片段缩放到规定的最大大小（ max_audio），并将其存储在数组 audio 中。然后，函数随机选择两个文件片段（即“语音抖动”），以用于数据增强。

对于每个片段，函数采用一个长度为 max_audio 的滑动窗口在语音片段中滑动，选择最具代表性的内容并将其保存在数组 feats 中。此外，函数确保两个片段不重叠，并将它们存储在大小为 [2, max_audio] 的 NumPy 数组 feat 中。最终，feat 被返回给 __getitem__(self, index) 函数，以用于神经网络的训练输入。

需要注意的是，代码中的处理方法可能与其他语音识别模型略有差异，因为数据预处理有很多途径，这里给出的只是其中一种可能性。