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)
函数,以用于神经网络的训练输入。需要注意的是,代码中的处理方法可能与其他语音识别模型略有差异,因为数据预处理有很多途径,这里给出的只是其中一种可能性。