小董的BLOG
0%

语音增强经典模型论文精读与实现

发表于 更新于

FullSubNet

核心思路：将一个纯全频带模型和一个纯子频带模型依次连接起来，并利用实际的联合训练将这两种模型的优点结合起来。

  • 全频带模型：输入全频带 带噪语音频谱，输出全频带 预测纯净语音 的模型。全频带模型可以捕获全局上下文谱和长距离交叉频带依赖，但缺乏信号平稳性建模和关注局部谱模式的能力
  • 子频带模型：模型的输入由一个频率和多个上下文频率组成。输出是对应频率的纯净语音。所有频率都是独立处理。由于噪声较于语音更加平稳，因此其可以通过对局部频带的平稳性建模区分噪声和语音；但对于信噪比极低的子带，效果不佳，因为其无法利用全频带信息。

通俗理解，子频带模型负责学习不同频段输入的分布，以得到不同频段的增强规则；全频带模型则负责学习全频带上频率的分布情况，以补充子频带模型没有利用全频带信息的缺陷

输入

首先需要强调一点，在FullSubNet中，两个模型的建模方向均为时间轴（核心！！！），也就是输出序列的长度为T

  • 全频带模型：输入直接是带噪语音的stft频谱
  • 子频带模型：全频带输出的格式为（B，1，F，T），原始FullSubNet会先将其reshape为（BxF，1，1，T），第一个1代表通道数，在此处没有意义；第二个1代表选择的当前频段，即全频带输出会将每个频段分离，每个频段信号都视作单个样本；此外，FullSubNet还会从原始带噪语音中选择相同频段和周围频段（（BxF，1，N，T），N为全频带当前样本频段的上下N个相邻频段），最后与全频带输出进行拼接，也就是子频带模型的输入为（BxF，1，1+N，T）

模型结构

全频带模型和子频带模型的结构完全一致，但由于子频带模型样本更小，因此其LSTM中隐藏层也相对小了一点；同时，由于子带模型输出的就是当前T时刻，对应频段的mask，因此也没有布置激活函数

原文使用的归一化是直接对每个样本的所有维度求均值再除以该均值

训练目标

模型的输出，也就是子频带模型输出为（B，F，2，T），分别代表实部虚部的mask，这样就可以同时修正幅度和相位：

而对于带噪语音，其可以表示为：

因此理想mask可以表示为：

然后直接拟合mask即可

总结

其实FullSubNet的核心就是在全频带模型后接了一个子频带模型，该子频带模型用共享的参数处理所有的子频带特征，其他就没啥了

DCCRN

CRN

在介绍DCCRN前，我们需要先知道它的核心框架，CRN：Convolutional Recurrent Network

在语音领域，CRN同时结合了CNN局部能力强、RNN能学习长程依赖的优点，提出了一个经典的Encoder-Decoder结构：

Encoder

由若干层Conv2D和BN、ELU组成

每层Conv2D的参数通常为：kernel=(5/3,2),stride=(2,1)对应输入形状（B，C，F，T），即Encoder会通过卷积逐步降采样频率信息，而时间长度则保持不变（但时间轴上也会做卷积）

对于离线增强任务，使用普通卷积，而对于在线增强推理任务，由于未来信息不可知，因此会使用掩码卷积（causal conv），由于卷积核是有宽度的，因此其原理会比Transformer中的causal mask使用的三角矩阵更复杂一丢丢：

假设现在数据为x0 x1 x2，卷积核大小为3，步长为1

那对x0的传统卷积窗口则为：pad x0 x1，即pad默认在两边

而在实时场景中t0时刻我们只有x0，没有x1，因此我们的causal conv的窗口即为：pad pad x0即全部在左边pad

再举个例子就是，我们如果对x2进行causal conv，窗口即为x0 x1 x2，在数值上等同于对x1进行普通conv，但是他们代表的意义是不同的（对应的时间步不同）

源码中的实现如下：

1
2
3
def forward(self,inputs):
    if self.padding[1] != 0 and self.causal:
        inputs = F.pad(inputs,[self.kernel_size-1, 0,0,0])

这里顺便说一下F.pad的用法，对于第二个参数，也就是这里的[self.kernel_size-1, 0,0,0]，其代表的是：

在最后一维的左边添加self.kernel_size-1个pad，右边不添加pad，在倒数第二维左右都不添加pad。最重要的就是其顺序是从最后一维开始指定，该参数最小为2维，最多不设上限，也就是这里也可以就是[self.kernel_size-1, 0]

LSTM

Encoder最终会输出（B，C’，F‘，T），CRN中将其reshape为（B，T，C‘ x F’）作为输入

1
2
3
4
# [2, 256, 4, 200] = [2, 1024, 200] => [2, 200, 1024]
lstm_in = e_5.reshape(batch_size, n_channels * n_f_bins, n_frame_size).permute(0, 2, 1)
lstm_out, _ = self.lstm_layer(lstm_in)  # [2, 200, 1024]
lstm_out = lstm_out.permute(0, 2, 1).reshape(batch_size, n_channels, n_f_bins, n_frame_size)  # [2, 256, 4, 200]

Decoder

Decoder使用ConvTranspose2d将LSTM输出还原为原始输入形状，对于反卷积，其原理就是插零+普通卷积进行上采样。为了防止上采样丢失高频细节，Decoder与Encoder之间还存在对应的skip connection以保留高频细节：

1
2
3
4
5
d_1 = self.tran_conv_block_1(torch.cat((lstm_out, e_5), 1))
d_2 = self.tran_conv_block_2(torch.cat((d_1, e_4), 1))
d_3 = self.tran_conv_block_3(torch.cat((d_2, e_3), 1))
d_4 = self.tran_conv_block_4(torch.cat((d_3, e_2), 1))
d_5 = self.tran_conv_block_5(torch.cat((d_4, e_1), 1))

最终，Decoder输出增强的Mask

CRN和DCCRN都是参考的U-Net的skip connection结构，即使用cat，而不是ResNet系的叠加，这样可以更好保留encoder和decoder各自的特征，但计算消耗会变大

CRN以及传统AE任务的缺陷

CRN预测的仅为幅值谱，在最终增强中使用的是noisy的相位，在低信噪比下性能有限。而对于预测复数谱的模型来说，输出的实部和虚部仅仅是输出的两个channel，模型并没有被显式告知这两个channel满足复数乘法规律，并且这些模型内部计算也全部都是实数表示，也就是说仅仅通过拟合来使模型最终输出倾向于一个实部值一个虚部值，并没有还原复数本身的结构

Complex Conv

因此，DCCRN直接提出了一个可以直接接收复数输入的卷积模块：

对于输入：

卷积核对应的也为：

按照复数乘法规律，则卷积输出为：（DCCRN核心公式）

也就是说，一层Complex Conv就进行了4次卷积，更符合STFT的物理意义

其具体代码就是将Wr和Wi分别定义为两个卷积层：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
if self.complex_axis == 0:#已将实部虚部cat在同一维
    real = self.real_conv(inputs)
    imag = self.imag_conv(inputs)
    real2real,imag2real = torch.chunk(real,2, self.complex_axis)
    real2imag,imag2imag = torch.chunk(imag,2, self.complex_axis)
    
else:
    if isinstance(inputs, torch.Tensor):
        real,imag = torch.chunk(inputs, 2, self.complex_axis)

    real2real = self.real_conv(real,)
    imag2imag = self.imag_conv(imag,)

    real2imag = self.imag_conv(real)
    imag2real = self.real_conv(imag)

real = real2real - imag2imag
imag = real2imag + imag2real
out = torch.cat([real, imag], self.complex_axis)

对应的，Batchnorm也加入了协方差来使复数整体归一化，激活函数也改为了复数输入形式

complex LSTM

格式与complex conv基本一致，也是将两层LSTM结构化复数格式，不再赘述

输出

DCCRN最终输出的是CRM（复数比例mask），即：

损失函数为复数谱的MSE和SI-SNR

LibriSpeech实战

模型参数配置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 创建模型 (默认配置)
model = DCCRN(
  rnn_layers=2,        # LSTM 层数
  rnn_units=128,       # LSTM 隐藏单元数
  win_len=400,         # STFT 窗长，16K对应25ms
  win_inc=100,         # STFT 帧移
  fft_len=512,         # FFT 点数
  win_type='hann',     # 窗函数类型
  masking_mode='E',    # 掩蔽模式: 'E', 'C', 'R'
  use_clstm=False,     # 是否使用复数LSTM
  use_cbn=False,       # 是否使用复数BatchNorm
  kernel_size=5,       # 卷积核大小
  kernel_num=[16,32,64,128,256,256]  # 各层通道数
)

对于masking_mode，对应网络最终输出的mask格式：

  • E：估计幅度掩蔽 + 相位修正 (tanh 限制幅度)
  • C：复数乘法掩蔽 (直接估计复数 mask)
  • R： 实数掩蔽 (实部和虚部分别乘 mask)

use_clstm和use_cbn则是选择是否要在LSTM和BN中也使用复数计算模型，当使用复数LSTM时，官方代码的卷积核对应也增大为了kernel_num=[32, 64, 128, 256, 256, 256]

在论文中，如下参数是性能最好的配置：

1
2
3
4
5
6
7
8
9
10
11
12
13
model = DCCRN(
    rnn_layers=2,
    rnn_units=128,
    win_len=400,
    win_inc=100,
    fft_len=512,
    win_type='hann',
    masking_mode='E',
    use_clstm=True,
    use_cbn=True,
    kernel_size=5,
    kernel_num=[32, 64, 128, 256, 256, 256],
)

前向传播

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def forward(self, inputs, lens=None):
    specs = self.stft(inputs)
    real = specs[:,:self.fft_len//2+1]
    imag = specs[:,self.fft_len//2+1:]
    spec_mags = torch.sqrt(real**2+imag**2+1e-8)
    spec_mags = spec_mags
    spec_phase = torch.atan2(imag, real)
    spec_phase = spec_phase
    cspecs = torch.stack([real,imag],1)
    cspecs = cspecs[:,:,1:]
    '''
    means = torch.mean(cspecs, [1,2,3], keepdim=True)
    std = torch.std(cspecs, [1,2,3], keepdim=True )
    normed_cspecs = (cspecs-means)/(std+1e-8)
    out = normed_cspecs
    ''' 

    out = cspecs
    encoder_out = []
    
    for idx, layer in enumerate(self.encoder):
        out = layer(out)
    #    print('encoder', out.size())
        encoder_out.append(out)
    
    batch_size, channels, dims, lengths = out.size()
    out = out.permute(3, 0, 1, 2)
    if self.use_clstm:
        r_rnn_in = out[:,:,:channels//2]
        i_rnn_in = out[:,:,channels//2:]
        r_rnn_in = torch.reshape(r_rnn_in, [lengths, batch_size, channels//2*dims])
        i_rnn_in = torch.reshape(i_rnn_in, [lengths, batch_size, channels//2*dims])
    
        r_rnn_in, i_rnn_in = self.enhance([r_rnn_in, i_rnn_in])

        r_rnn_in = torch.reshape(r_rnn_in, [lengths, batch_size, channels//2, dims])
        i_rnn_in = torch.reshape(i_rnn_in, [lengths, batch_size, channels//2, dims]) 
        out = torch.cat([r_rnn_in, i_rnn_in],2)
    
    else:
        # to [L, B, C, D]
        out = torch.reshape(out, [lengths, batch_size, channels*dims])
        out, _ = self.enhance(out)
        out = self.tranform(out)
        out = torch.reshape(out, [lengths, batch_size, channels, dims])
   
    out = out.permute(1, 2, 3, 0)
    
    for idx in range(len(self.decoder)):
        out = complex_cat([out,encoder_out[-1 - idx]],1)
        out = self.decoder[idx](out)
        out = out[...,1:]
    #    print('decoder', out.size())
    mask_real = out[:,0]
    mask_imag = out[:,1] 
    mask_real = F.pad(mask_real, [0,0,1,0])
    mask_imag = F.pad(mask_imag, [0,0,1,0])
    
    if self.masking_mode == 'E' :
        mask_mags = (mask_real**2+mask_imag**2)**0.5
        real_phase = mask_real/(mask_mags+1e-8)
        imag_phase = mask_imag/(mask_mags+1e-8)
        mask_phase = torch.atan2(
                        imag_phase,
                        real_phase
                    ) 

        #mask_mags = torch.clamp_(mask_mags,0,100) 
        mask_mags = torch.tanh(mask_mags)
        est_mags = mask_mags*spec_mags 
        est_phase = spec_phase + mask_phase
        real = est_mags*torch.cos(est_phase)
        imag = est_mags*torch.sin(est_phase) 
    elif self.masking_mode == 'C':
        real,imag = real*mask_real-imag*mask_imag, real*mask_imag+imag*mask_real
    elif self.masking_mode == 'R':
        real, imag = real*mask_real, imag*mask_imag 
    
    out_spec = torch.cat([real, imag], 1) 
    out_wav = self.istft(out_spec)
     
    out_wav = torch.squeeze(out_wav, 1)
    #out_wav = torch.tanh(out_wav)
    out_wav = torch.clamp_(out_wav,-1,1)
    return out_spec,  out_wav

ConvSTFT

在DCCRN中，由于2020年torch还没有推出支持自动求导的STFT，因此其使用Conv1d函数来实现了一个可微的STFT，需要注意的是，此处使用Conv1d正常来说仅是为了实现STFT，因此其权重是固定为傅里叶基的

由代码可得知，模型的输入直接就是（b，采样点），因此直接输入音频即可

另外还有一点要说明一下，在语音领域一个常用的STFT配置是：

1
2
3
4
fs=16000
win_length=400
hop=100
n_fft=512

此处窗长小于了傅里叶变换次数，但对于离散傅里叶变换（可见信号处理基础算法 | 小董的BLOG），我们要求时间采样点（也就是STFT的窗长）要等于傅里叶变换次数，在几乎所有库函数中，会自动补0，使样本长度=512。像这样配置的好处是可以用更少的时域数据获得更平滑的频率轴，同时也可以增大输入维度，但这实际上不能真正增加频域分辨率

加载数据集

代码比较繁长，就不放了，注意以下几点即可

噪声信噪比范围问题

语音使用librispeech，噪声使用选煤厂噪声，当时测下来噪声平均声压级为100 db左右，如果假设大声说话声压级为80 db左右，因此最大信噪比会来到<-20 db，实际训练信噪比取[-20,10]

音频长度截取

DCCRN和大部分离线训练的AE模型很多都是取4s为一个样本。对于每个音频，<4s的进行重复拼接，>4s的则随机裁剪4s的长度

归一化问题

如果是直接加载的wav，那其默认数值范围就是[-1,1]。由于信号混合时都是直接通过信噪比控制数值，而通常做法就是根据语音的能量（RMS）控制噪声的增益来实现信噪比的控制，所以要先将语音控制在[-1,1]，通常有两种做法来实现归一化：

  • 由于wav加载的数值都是[-1,1]（如果不是就直接归一化），因此直接对带噪信号进行峰值归一化：

    1
    2
    3
    4
    5
    6
    noisy = clean + noise
    # 峰值归一化: 混合后可能超出 [-1, 1], 同步缩放 noisy 和 clean
    peak = np.max(np.abs(noisy)) + 1e-10
    if peak > 1.0:
        noisy = noisy / peak
        clean = clean / peak

    即仅对较强信号进行一个峰值缩放，注意语音信号也要进行相同缩放，这样才能使语音本身数值一致

  • 直接控制clean和noise本身的能量RMS，先对齐到同一较低级别，这样即使相加也不会超过[-1,1]的范围，且可以确保整个数据集中所有样本的能量分布符合真实的声学统计规律，而不是被强制拉伸到统一的峰值。

另外，为了防止模型输出音频也超过范围，通常会在模型最后添加一个Tanh激活函数（直接生成波形）或者直接映射到sigmoid（mask方法）

损失函数

源论文是直接使用的仅SI-SDR，也可以使用MSE等常规loss

1
2
3
4
5
6
7
8
9
10
11
def si_snr(s1, s2, eps=1e-8):
    #s1 = remove_dc(s1)
    #s2 = remove_dc(s2)
    s1_s2_norm = l2_norm(s1, s2)
    s2_s2_norm = l2_norm(s2, s2)
    s_target =  s1_s2_norm/(s2_s2_norm+eps)*s2
    e_nosie = s1 - s_target #分母
    target_norm = l2_norm(s_target, s_target)#分子
    noise_norm = l2_norm(e_nosie, e_nosie)
    snr = 10*torch.log10((target_norm)/(noise_norm+eps)+eps)
    return torch.mean(snr)

SI-SDR值越大越好，因此作为loss的话要取一个负号

离线训练转在线流式推理

流式缓存维护

LSTM缓存维护

由于在流式推理中，是把很多chunk拼在一起，而网络是一个chunk一个chunk的处理数据，为了保证处理当前chunk时保留之前chunk的记忆，需要对模型中LSTM进行缓存维护：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def forward(self, inputs, states=None):
    if isinstance(inputs,list):
        real, imag = inputs
    elif isinstance(inputs, torch.Tensor):
        real, imag = torch.chunk(inputs,-1)

    # states format:
    #   ( (h_r, c_r), (h_i, c_i) )  for real_lstm
    #   ( (h_r, c_r), (h_i, c_i) )  for imag_lstm
    if states is not None:
        (h0_rl_r, c0_rl_r), (h0_rl_i, c0_rl_i) = states[0]  # real_lstm
        (h0_il_r, c0_il_r), (h0_il_i, c0_il_i) = states[1]  # imag_lstm
    else:
        h0_rl_r = c0_rl_r = h0_rl_i = c0_rl_i = None
        h0_il_r = c0_il_r = h0_il_i = c0_il_i = None

    # real_lstm processes 'real' and 'imag' independently
    if h0_rl_r is not None:
        r2r_out, (h_rl_r, c_rl_r) = self.real_lstm(real, (h0_rl_r, c0_rl_r))
        i2r_out, (h_rl_i, c_rl_i) = self.real_lstm(imag, (h0_rl_i, c0_rl_i))
    else:
        r2r_out, (h_rl_r, c_rl_r) = self.real_lstm(real)
        i2r_out, (h_rl_i, c_rl_i) = self.real_lstm(imag)

    # imag_lstm processes 'real' and 'imag' independently
    if h0_il_r is not None:
        r2i_out, (h_il_r, c_il_r) = self.imag_lstm(real, (h0_il_r, c0_il_r))
        i2i_out, (h_il_i, c_il_i) = self.imag_lstm(imag, (h0_il_i, c0_il_i))
    else:
        r2i_out, (h_il_r, c_il_r) = self.imag_lstm(real)
        i2i_out, (h_il_i, c_il_i) = self.imag_lstm(imag)

    real_out = r2r_out - i2i_out
    imag_out = i2r_out + r2i_out
    if self.projection_dim is not None:
        real_out = self.r_trans(real_out)
        imag_out = self.i_trans(imag_out)

    new_states = (
        ((h_rl_r, c_rl_r), (h_rl_i, c_rl_i)),   # real_lstm
        ((h_il_r, c_il_r), (h_il_i, c_il_i)),   # imag_lstm
    )
    return [real_out, imag_out], new_states

源码中，共有2个LSTM，这2个LSTM会分别都进行一次实数和虚数计算，而这些计算都应该分开维护，因此最终我们需要维护的缓存共有4套，即对于同一LSTM，实数计算和虚数计算的隐藏层和细胞层都应该独立维护

nn.LSTM的输入可以直接在第二个参数输入隐藏层和细胞层参数（即上一chunk的最后一个时间步的隐藏状态），其输出固定为[output,(h, c)]

卷积层维护

当模型逐个chunk处理时，为了实现causal conv，每个chunk在左侧都会进行1个pad（卷积核大小为2的情况），所以理论上每个chunk在第每一层encoder，decoder中都有一个开头0pad，那我们在后续chunk的建模中，需不需要把这个0pad替换为上一chunk的最后一个时间步呢？具体答案会在下一节流式数据处理中讲到

流式数据处理

对于流式数据的处理，有两种方法，它们的核心就是是否使用额外overlap

首先，无论是哪种方法，chunk之间都一定会存在overlap，它们是为了补充以下内容：

  • STFT的padding。除非STFT使用center=false，所有的STFT会在左、右侧添加padding（DCCRN自定义的STFT只填充左边），防止对首尾加窗时超出界限。对于第一个之后的chunk，我们希望这个padding由上一段音频替代，这样可以获得更稳定的首帧数据，而我们如果将上一chunk尾部的部分数据拼接到当前chunk，就可以实现。如果overlap的长度=这个stft的padding，那么每个chunk之间的STFT频谱边界就是完全连续的

    其实最好的做法就是使用center=false，这样甚至都不用overlap了，但这样在训练中就会麻烦一点点。

  • causal Conv的padding。前面我们提到过causal conv的实现也是通过padding，我们同样不希望除第一个chunk之外的chunk出现padding，这样每一个chunk在卷积层都会被视作是“第一个chunk”。因此对于每一个causal conv，我们都需要上一个chunk的最后一个时间帧的特征。

    举个例子，DCRNN由6层encoder和decoder组成，每层中都有一个causal conv，其会在输入特征第一个时间帧之前做一帧padding，因此对于第一层causal conv，我们需要上一个chunk对应的第一层输入的最后一帧来代替padding，卷积会将这层padding帧和第一个特征帧合为一帧，因此在第二层causal conv，我们就又需要上一个chunk的第二层输入的最后一帧，注意层与层之间的缓存是一一对应的。

​ 通过如上分析，我们想实现conv层的缓存就有两种方法：

  • 维护每一层conv输入的最后一帧缓存供下一个chunk使用，这样就不需要用额外overlap了，这样效率最高，缺点就是每一层都要维护，稍微麻烦一点

  • 加入额外的overlap，让STFT重合的时间帧完全覆盖causal conv的padding范围。对于第一层conv输入，其需要上一chunk的最后一帧，对于第二层，则需要上一chunk的原本倒数第二帧（因为倒一已经被上一层conv融合了），这样算下来，标准DCCRN则需要大概13帧STFT时间帧的额外overlap（大概就是100ms），转换到时域上，这段overlap并不算短，并且这就要求当前chunk要大于100ms，这会引入额外的延迟。优点就是这样就不用专门去维护conv内存了，比较简单暴力

    再额外提一嘴，LSTM的缓存是必须维护的，因为其包括了之前所有时间帧记忆

根据上面两种方法，就可以得到两种流式数据处理的方法：

  • 小chunk法，仅考虑STFT的overlap，这样我们每个chunk就可以设的很小（极端地，可以直接把每个时间帧都当做一个chunk，也就是说直接对每个窗长的时域数据做FFT），STFT的overlap长度大概需要12.5ms（400*0.5/16000），我们就可以把chunk定为100甚至50ms以内，可以明显降低延迟

    这样操作需要自定义一个STFT，因为padding需要自己替换。当使用torch的stft，center=false时，就不需要考虑overlap，直接拼接各chunk即可

  • 大chunk法，通过大量overlap避免直接维护conv缓存，这样我们的overlap大概会来到100ms以上，因此chunk大小也会跟着变大，从而使延迟增加

这样看下来，小chunk法确实更好，如果我们有效chunk定为50ms，那除了第一个chunk，其他chunk实际长度则为50+12.5=62.5ms（或者说第一个chunk的12.5ms全为0padding），多出来的12.5ms会在转stft中消失，最终模型输出仍为50ms