Conv1D Operations

今天来总结一下在我过去的一些任务上用到的一些一维卷积技巧，主要是depthwise和transposed conv1d。之所以写一篇文章是觉得存在两方面的好处，其一是可以加深对现代神经网络架构中卷积操作逻辑的理解，二是能提升编码和运算上的效率，让代码本身更加简洁。

FSMN & Beamformer

FSMN之前简单的写过一篇文章做介绍，是Shiliang Zhang在2016年提出的一种无反馈结构的，可以用于长序列建模的网络类型，类似的还有TCN（TDNN的延伸）等等。在第$\ell$层，$t$时刻的输出$\mathbf{h}_{t, \ell}$由如下运算产生：

$\mathbf{h}_{t, \ell} = f(\mathbf{W}_{\ell} \mathbf{h}_{t, \ell - 1} + \tilde{\mathbf{W}}_{\ell} \tilde{\mathbf{h}}_{t, \ell - 1} + \mathbf{b}) \tag{1}$

其中$\tilde{\mathbf{h}}_{t, \ell - 1}$被称为memory blocks，用于建模上下文信息，在vectorized版本中，表示成如下加权的形式：

$\tilde{\mathbf{h}}_{t, \ell - 1} = \sum_{c = 0}^{C} \mathbf{a}_{c, \ell - 1} \odot \mathbf{h}_{t + c - N_l - 1, \ell - 1} \tag{2}$

其中$C = N_r + N_l + 1，N_l, N_r$分别表示左右context的长度。此式的计算可以理解为一个depthwise的一维卷积过程，其中卷积核的大小为$C$，步长为1（原始版本不做下采样），输出通道数为$D$，其中$D$表示$\mathbf{a}_{i, \ell - 1}$和$\mathbf{h}_{i, \ell - 1}$的维度。代码上下面的定义进行卷积：

# definition
ctx_conv = nn.Conv1d(memory_size,
                     memory_size,
                     kernel_size=ctx_size,
                     dilation=dilation,
                     groups=project_size,
                     padding=(ctx_size - 1) // 2,
                     bias=False)
# call
# h: batch x num_frames x memory_size
# m: batch x num_frames x memory_size
m = self.ctx_conv(h)

depthwise conv1d的一个非常常见的应用就是用于替代一组线性层，代码上看起来显得更加简洁，比如multi-head attention，beamformer和这里的FSMN中，都可以加以运用。pytorch的api里面，选取groups的值（通常等于输入通道数）来隔绝输入的通道之间的信息交换（分组进行卷积操作并输出到对应的通道中）即可。

还有一点需要注意的是，虽然结果（数学）上是完全等价的，但是逻辑上， conv1d的操作视角不是以时间为单元（也就是语音帧），而是通道，对应特征的维度。比如在上面的ctx_conv层中，卷积核$\mathbf{K} \in \mathbf{R}^{C \times D}$在$c$个通道上学到的结果对应为$\mathbf{k}_c = [a_{c, i}, \cdots, a_{c, N_r + N_l}]$，跨越时间而通道固定。

拓展一下可以用于实现一组（固定）波束形成器。信号上我们一般这样定义一个指向为$\theta$的波束形成器，即

$s_{t,f}^\theta = \mathbf{w}_{\theta, f}^H \mathbf{y}_{t,f} \tag{3}$

其中$\mathbf{w}_{\theta, f}, \mathbf{y}_{t,f} = [y_{0,t,f}, \cdots, y_{M - 1,t,f}] \in \mathbf{C}^{M}$，$M$为通道数。从上式可以看出，beamformer的权重在每个频率上是独立的，因此可以用depthwise conv1d进行实现，加上操作在复数域，再用两个conv1d定义一层复数卷积即可。式$(3)$可用下面的卷积层进行计算，需要注意的是，由于每一帧的操作逻辑相同，因此需要将输入的时间维度和batch的维度（通常也就是第一维）合并：

# definition
beamformer = ComplexConv1d(num_bins,
                           num_bins,
                           num_channels,
                           groups=num_bins,
                           padding=0,
                           bias=False)
# call
# x: batch*num_frames x num_bins x num_channels, complex tensor
# beam: batch*num_frames x num_bins x 1, complex tensor
beam = beamformer(x)

如果是$\Theta$个波束形成器，只需要将输出通道数调整为$\Theta \times F$即可，对应的输出则为batch*num_frames x num_beams*num_bins x 1。

最后小结一下，如果想用depthwise conv1d去做一些运算的简化或者实现，需要在原表达式中找到通道维度上的线性变换操作，并且满足一定的独立性要求。合适的使用可以将原本看起来比较复杂的表达式（加和，点乘，线性变换等等）用一个卷积操作完成，代码上精炼很多。

(i)STFT

本部分说一下用（转置） conv1d做(i)STFT的逻辑，虽然很早之前就在用了，但是实现上的细节以及背后的原理存在有一些需要解释的地方。信号上的(i)STFT分别对应分帧，加窗，DFT以及iDFT，加窗，重叠相加（Overlap and Add，OLA）的过程，其中分帧和OLA可以借用conv1d和transposed conv1d来进行，同时将卷积核初始化为DFT变换矩阵和窗函数的乘积，以合并加窗和(i)DFT的逻辑。

用$\mathbf{K}_{\text{(i)DFT}} \in \mathbf{C}^{N \ \times N}$表示(i)DFT变换矩阵，满足如下正交关系：

$\mathbf{K}_{\text{DFT}}^H \mathbf{K}_{\text{iDFT}} = I \tag{4}$

一般的，为了满足这个条件，有两种可以选择的初始化方法，分别为：

$\begin{cases} \mathbf{K}_{\text{DFT}} = \mathbf{W}/\sqrt{N}, \mathbf{K}_{\text{iDFT}} = \mathbf{W}^H/\sqrt{N} \\ \mathbf{K}_{\text{DFT}} = \mathbf{W}, \mathbf{K}_{\text{iDFT}} = \mathbf{W}^H/N \end{cases} \tag{5}$

其中$\mathbf{W}$表示单位矩阵的DFT变换结果，即$\mathcal{F}(\mathbf{I})$，$N$为FFT bins的个数。第一种中，正逆变换矩阵为酉矩阵，第二种一般是我们写(i)FFT的逻辑，对iFFT的结果需要除一个$N$，相当于在变换矩阵上加一个同等的系数。语音信号为实信号，所以STFT拿到的实部和虚部可看成是卷积核的实部和虚部单独作用的结果，表示为：

$\begin{align} \mathbf{Y}^R &= \text{conv}_{\text{1D}} \left(\mathbf{x}, \mathbf{K}_{\text{DFT}}^R \odot \mathbf{w}, S\right) \notag \\ \mathbf{Y}^I &= \text{conv}_{\text{1D}} \left(\mathbf{x}, \mathbf{K}_{\text{DFT}}^I \odot \mathbf{w}, S \right) \notag \end{align} \tag{6}$

其中$\mathbf{x}$为语音信号序列，$\mathbf{w}$表示STFT的分析窗，$S$为帧移。合并成一个卷积操作后即：

$\text{STFT}(\mathbf{x}) = [\mathbf{Y}^R; \mathbf{Y}^I]^T = \text{conv}_{\text{1D}} \left(\mathbf{x}, \left[\mathbf{K}_{\text{DFT}}^R; \mathbf{K}_{\text{DFT}}^I \right]^T \odot \mathbf{w}, S\right) \tag{7}$

$\mathbf{K} = [\mathbf{K}_{\text{DFT}}^R; \mathbf{K}_{\text{DFT}}^I]^T \in \mathbf{R}^{2N \times N}$。Pytorch中可以对应的初始化为（实际需要考虑窗长和FFT bins数不同的情况）：

def init_kernel(frame_len,
                frame_hop,
                window,
                round_pow_of_two=True,
                normalized=False,
                inverse=False):
    """
    Return STFT kernels
    """
    # FFT points
    B = 2**math.ceil(math.log2(frame_len)) if round_pow_of_two else frame_len
    if normalized:
        # make K^H * K = I
        S = B**0.5
    else:
        S = 1
    I = th.stack([th.eye(B), th.zeros(B, B)], dim=-1)
    # W x B x 2
    K = th.fft(I / S, 1)[:frame_len]
    if inverse and not normalized:
        # to make K^H * K = I
        K = K / B
    # 2 x B x W
    K = th.transpose(K, 0, 2) * window
    # 2B x 1 x W
    K = th.reshape(K, (B * 2, 1, frame_len))
    return K

需要对上面的代码做几点解释。第一，为什么要卷积核选择成全DFT的结果而不是rFFT的结果？因为如果只是用于一般声学特征的提取，取K = th.rfft(th.eye(B) / S, 1)[:frame_len]即可。这是考虑到，在做iSTFT的时候我们需要一个全DFT的核。第二，为什么没有给K取逆和共轭。不取逆是因为DFT的矩阵本身是酉矩阵，共轭转置就是自身的逆，因此，在inverse的情况下，返回的卷积核跟实际iDFT使用的其实只差了一个共轭而已，至于这个共轭为何不取，下面的文章会进行解释。

在init_kernel的初始化kernel作用下，返回的STFT结果是一个full FFT的形式，对于一般声学特征的提取，取前$N / 2 + 1$部分就行。若将$[\mathbf{Y}^R; \mathbf{Y}^I]^T $变换到时域，使用一维转置卷积表示为：

$\begin{align} \mathbf{x}' & = \text{deconv}_{\text{1D}}([\mathbf{Y}^R; \mathbf{Y}^I]^T, \left[\mathbf{K}_{\text{DFT}}^R; \mathbf{K}_{\text{DFT}}^I \right]^T \odot \mathbf{w}, S) \notag \\ & = \text{deconv}_{\text{1D}}(\mathbf{Y}^R, \mathbf{K}_{\text{DFT}}^R \odot \mathbf{w}, S) + \text{deconv}_{\text{1D}}(\mathbf{Y}^I, \mathbf{K}_{\text{DFT}}^I \odot \mathbf{w}, S) \notag \end{align} \tag{8}$

这么做背后的原理有二，其一是iSTFT中，对于每一帧，我们实际只使用iDCT变换的实部，根据复数运算$(a + bi)(c + di) = (ac - bd) + (ad + bc)i$，只需要保留$ac - bd$的实部即可。第二问题就是，$(8)$式中连接实部结果和虚部结果的符号是加号，结果依然正确，如何解释？。因为我们实际使用的卷积核是没有取共轭的结果，所以真实的结果应该是$(a + bi)(c - di) = (ac + bd) + (ad - bc)i$中的$ac + bd$，负号变成加号，和$(8)$式子的结果一致。分析之后，可以简单验证一下一帧变换的结果：

B = 8
th.random.manual_seed(666)
w = init_window("rect", B)
W = init_kernel(B, B // 2, w, normalized=True)
x = th.rand(1, 1, B * 2)
x[..., B:] = 0
p = tf.conv1d(x, W, stride=B, padding=0)
y = tf.conv_transpose1d(p, W, stride=B, padding=0)
print(f"x = {x.squeeze()}")
print(f"y = {y.squeeze()}")
print(f"error: {th.sum(x - y).item():f}")

验证输出

x = tensor([0.3119, 0.2701, 0.1118, 0.1012, 0.1877, 0.0181, 0.3317, 0.0846, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000])
y = tensor([0.3119, 0.2701, 0.1118, 0.1012, 0.1877, 0.0181, 0.3317, 0.0846, 0.0000,
        0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000])
error: 0.000000

如果不清楚一维转置卷积的逻辑的同学可以写几个toy验证一下overlap and add的结果，这样就能理解到本部分的一些符号表达了。网上有一些关于通用转置卷积的分析，不过理解起来有些麻烦，在一维条件下，从这个角度切入就清晰很多。最后关于iSTFT的实现还有几个细节问题，解释如下：

在做iSTFT的时候，如果输入为onesided的STFT结果（即前$N / 2 + 1$个），需要先padding到$N$个，之后参与到转置卷积中去。具体的逻辑就是根据DFT结果的对称性，实部关于$N / 2 + 1$对称，实部共轭对称。

主流的iSTFT会在时域的信号上加一个根据窗函数计算的归一化系数，用于保证变换前后的幅度一致（以及解决窗本身的一些正交性问题），比如librosa里面的，和torchaudio里面的，逻辑就是对窗函数的能量也做一个OLA的过程，因此同样可以用转置卷积得到：

win = th.repeat_interleave(window[None, ..., None],
                           packed.shape[-1],
                           dim=-1)
# W x 1 x W
I = th.eye(window.shape[0], device=win.device)[:, None]
# 1 x 1 x T
norm = tf.conv_transpose1d(win**2, I, stride=frame_hop, padding=0)

全部分的code由于篇幅我没有放在文章中，按照本文的解释是可以完整无误的写出来的，之后可以测几个音频样例，算一下误差验证算法逻辑。

Conclusion

最后小结一下吧。这部分内容其实是过去几年经常用到的conv1d操作，乘着最近在做ASR的repo，便来小结一下，关于(i)STFT那部分花了一些时间去理解背后的逻辑和设计问题，早期网上有一些版本的实现，但是个人觉得多少有些问题，比如操作上有偏差或者冗余等等。类似的操作在实际编码中个人觉得还是需要注意的，因为知道每种操作背后的逻辑之后，可以辅助的产生一些高效的编码，一定程度上也可以节省工程或者实验的时间。