SUDORMRF: EFFICIENT NETWORKS FOR UNIVERSAL AUDIO SOURCE SEPARATION

人的精神寄托可以是音乐，可以是书籍，可以是运动，可以是工作，可以是山川湖海，唯独不可以是人。

Depthwise Separable Convolution 深度分离卷积（前置知识）

相比于常规卷积，可以减少参数和计算量。分为两步：

1. Depthwise Convolution (DW) 深度卷积：每个卷积核只处理一个输入通道（卷积核深度均为1）。因此输入特征的通道数不变。
2. Pointwise Convolution (PW) 逐点卷积：卷积核尺寸为$1\times 1$，输入特征的尺寸不变。

我们以input feature map shape为$[12,12,3]\in R^{H\times W\times C}$，卷积核shape为$[5,5,3]\in R^{k_h\times k_w\times C}$, stride = 1, padding = 0，卷积核数量为256（输出通道数为256），output feature map shape为$[8,8,256]$的情况为例。

• 常规卷积中，卷积核的参数量：$C_{\text{in}}\times C_{\text{out}}\times k_{\text{h}}\times k_{\text{w}}$
  计算量（FlOPs）：$C_{\text{in}}\times C_{\text{out}}\times k_{\text{h}}\times k_{\text{w}}\times O_{\text{h}}\times O_{\text{w}}$
  其中$O_{\text{h}}\times O_{\text{w}}$为output feature map的高$\times$宽，$k_{\text{h}}\times k_{\text{w}}$为卷积核对应尺寸。（对于PW，DW而言，计算公式稍有变化）

1. 采用常规卷积：

• 参数量：$3\times 256\times 5\times 5=19200$
• FLOPs：$3\times 256\times 5\times 5\times 8\times 8=1228800$

2. 采用Depthwise Separable Convolution:
   将$[5,5,3]\times 256$的卷积层拆分为两个卷积层，先后经过两层卷积以完成同样的维度转换效果:
   1. DW: $[5,5,1]\times 3$
   2. PW: $[1,1,3]\times 256$

• 参数量：DW： $1\times 3\times 5\times 5$（这里不是$C_{\text{in}}$，因为每个卷积核只负责一个通道，即卷积核深度为1）；PW： $3\times 256\times 1\times 1$。总计为：843
• FLOPs: $1\times 3\times 5\times 5\times 8\times 8+3\times 256\times 1\times 1\times 8\times 8=53952$。（分两步卷积，先DW后PW）

2.Methodology

Overall Architecture

整体算法流程如下：

定义输入输出尺寸
$\mathbf{x}\in {\mathbb{R}}^T$为混合音频信号，$\mathcal{E}$为Encoder，对输入$\mathbf{x}$处理得到特征向量: ${\mathbf{v}}_{\mathbf{x}}=\mathcal{E}\left(\mathbf{x}\right)\in {\mathbb{R}}^{C_{\epsilon }\times L}$。将${\mathbf{v}}_{\mathbf{x}}$送入Separation Module $S$当中得到${\hat{\mathbf{m}}}_i\in {\mathbb{R}}^{C_{\epsilon }\times L}$为第$i$个音源的mask，$i=1,2,\cdots ,N$。假设共有$N$个音源产生的音频信号${\mathbf{s}}_1,{\mathbf{s}}_2,\cdots ,{\mathbf{s}}_N\in {\mathbb{R}}^T$共同组成$\mathbf{x}$。将${\mathbf{v}}_{\mathbf{x}}$与$\widehat{{\mathbf{m}}_i}$逐项相乘得到第$i$音源的特征向量$\widehat{{\mathbf{v}}_i}$：
$$\widehat{{\mathbf{v}}_i}={\mathbf{v}}_{\mathbf{x}}\odot \widehat{{\mathbf{m}}_i}$$
再经过解码器$\mathcal{D}$得到${\hat{\mathbf{s}}}_i=\mathcal{D}(\widehat{{\mathbf{v}}_i})$

• 定义1：${\text{Conv1D}}_{C,K,S}:{\mathbb{R}}^{C_{in}\times L_{in}}\to {\mathbb{R}}^{C\times L}$。表示一维常规卷积。将输入shape从${\mathbb{R}}^{C_{in}\times L_{in}}$转为${\mathbb{R}}^{C\times L}$。其中$C$为output channel，$S$为stride，$K$为kernel size，L为尺度（时间）。

• 定义2：${\text{ConvTr1D}}_{C,K,S}:{\mathbb{R}}^{C_{in}\times L_{in}}\to {\mathbb{R}}^{C\times L}$。转置卷积。需要注意的一点是，转置卷积的运算过程，相当于原卷积核对input求梯度（具体过程参看参考链接）。

• 定义3：${\text{DWConv1D}}_{C,K,S}:{\mathbb{R}}^{C_{in}\times L_{in}}\to {\mathbb{R}}^{C\times L}$。一维Depthwise Convolution深度卷积。
  本质是将原始卷积层拆分成 $G=C_{\text{in}}$ 个Conv1D：${\hat{\mathcal{F}}}_i=[{\text{Conv1D}}_{C_G,K,S}{]}_i$，其中 i ∈ { 1 , ⋯   , G } ， C G = [ C / G ] i\in \{ 1, \cdots , G\}，C_G = [C / G] i∈{1,⋯,G}，CG​=[C/G]。每一个卷积核$\hat{\mathcal{F}}$对输出贡献$C_G$个通道。最终将 { F ^ i ∣ i = 1 , 2 , ⋯   , G } \{\hat{\mathcal{F}}_i | i = 1,2,\cdots,G \} {F^i​∣i=1,2,⋯,G}的输出结果按通道拼接：
  D W C o n v 1 D C , K , S ( x ) = C o n c a t ( { F i ( x i ) , ∀ i } ) , (2) \mathrm{DWConv}1\mathrm{D}_{C,K,S}\left(\mathbf{x}\right)=\mathrm{Concat}\left(\left\{\mathcal{F}_i\left(\mathbf{x}_i\right), \forall i\right\}\right),\tag2 DWConv1DC,K,S​(x)=Concat({Fi​(xi​),∀i}),(2)
  $\text{Concat}(\cdot )$表示拼接。

2.1 Encoder

Encoder用$\mathcal{E}$表示，包含一个一维卷积，kernel size为$K_{\mathcal{E}}$, stride为$K_{\mathcal{E}}/2$，使用公式表示Encoder的具体操作如下：
$${\mathbf{v}}_{\mathbf{x}}=\mathcal{E}\left(\mathbf{x}\right)=\mathrm{ReLU}\left(\mathrm{Conv}1{\mathrm{D}}_{C_{\mathcal{E}},K_{\mathcal{E}},K_{\mathcal{E}/2}}\left(\mathbf{x}\right)\right)\in {\mathbb{R}}^{C_{\mathcal{E}}\times L}\text{\hspace{1em}(3)}$$

其中ReLU为逐项激活，$C_{\mathcal{E}}$为Encoder的输出通道数。

2.2 分离模块

分离模块$S$将${\mathbf{v}}_x$做以下处理

1. 使用LN和Pointwise Conv将${\mathbf{v}}_x$映射到新的通道空间当中：
   $${\mathbf{y}}_0=\mathrm{Conv}1{\mathrm{D}}_{C,1,1}\left(\mathrm{LN}\left({\mathbf{v}}_{\mathbf{x}}\right)\right)\in {\mathbb{R}}^{C\times L}\text{\hspace{1em}(4)}$$

   $\text{LN}({\mathbf{v}}_{\mathbf{x}})$为layer-norm layer。

2. 使用多个BU-convolutional blocks（U-ConvBlocks）拼接在一起。其中第$i$个Block的输出作为第$i+1$个Block的输入。U-ConvBlock的具体细节参考Sec2.2.1，类似TDA-Net和U-Net，利用了多尺度信息，Block输入输出尺度一致。

3. 使用最后的Block输出${\mathbf{y}}_B^T\in {\mathbb{R}}^{L\times C}$，针对每一个音源设置一个Conv1D层，以得到对应音源的中间特征向量${\mathbf{z}}_i$（比如有N个音源，则对应N个Conv1D，随后使用${\mathbf{z}}_i$获得$\widehat{{\mathbf{m}}_i}$）：
   $${\mathbf{z}}_i=\text{Conv}1{\text{D}}_{C,C_{\mathcal{E}},1}{\left({\mathbf{y}}_B^T\right)}^T\in {\mathbb{R}}^{C_{\mathcal{E}}\times L}\text{\hspace{1em}(5)}$$
   其中${\mathbf{y}}_B^T$表示${\mathbf{y}}_B$的转置。

4. 使用${\mathbf{z}}_i$ 求解$\widehat{{\mathbf{m}}_i}$：
   $${\hat{\mathbf{m}}}_i={\mathrm{vec}}^{-1}\left(\frac{\exp \left(\mathrm{vec}\left({\mathbf{z}}_{\mathrm{i}}\right)\right)}{{\sum }_{j=1}^N\exp \left(\mathrm{vec}\left({\mathbf{z}}_j\right)\right)}\right)\in {\mathbb{R}}^{C_{\mathcal{E}}\times L}\text{\hspace{1em}(6)}$$

其中${\hat{\mathbf{m}}}_i\in [0,1{]}^{C_{\mathcal{E}}\times L}$。 vec$\left(\cdot \right):{\mathbb{R}}^{K\times N}\to {\mathbb{R}}^{K\cdot N}$，表示向量化。${\mathrm{vec}}^{-1}\left(\cdot \right):{\mathbb{R}}^{K\cdot N}\to$ ${\mathbb{R}}^{K\times N}$表示反向量化。可以发现，对于所有mask的同一位置，求和为1：
$$\sum _i^N{\widehat{{\mathbf{m}}_{\mathbf{i}}}}_{(x,y)}=1$$
其中$(x,y)$表示mask的某个元素的坐标。

5. 利用${\hat{\mathbf{m}}}_i$以及特征向量${\mathbf{v}}_x$得到每个音源的特征向量${\hat{\mathbf{v}}}_i$：
   $${\hat{\mathbf{v}}}_i={\mathbf{v}}_{\mathbf{x}}\odot {\hat{\mathbf{m}}}_i\in {\mathbb{R}}^{C_{\mathcal{E}}\times L}\text{\hspace{1em}(7)}$$

2.2.1 U-ConvBlock

U-ConvBlock整体架构与算法流程如下图所示：

整体而言与U-Net类似，但与TDA-Net更相似，也是连续的下采样，上采样。

• 定义4：${\mathrm{PReLU}}_C:{\mathbb{R}}^{C\times L}\to {\mathbb{R}}^{C\times L}$。（parametric
  rectified linear unit）：
  $${\mathrm{PReLU}}_C{\left(\mathbf{y}\right)}_{i,j}=\max \left(0,{\mathbf{y}}_{i,j}\right)+{\mathbf{a}}_i\cdot \min \left(0,{\mathbf{y}}_{i,j}\right)\text{\hspace{1em}(8)}$$
  ${\mathbf{a}}_i$为可学习的参数，$\mathbf{y}$为输入。

• 定义5：${\mathcal{I}}_M:{\mathbb{R}}^{C\times L}\to {\mathbb{R}}^{C\times M\cdot L}$。上采样操作，最邻近插值，$M$为缩放系数。

2.3 Decoder

$\mathcal{D}$表示Decoder，将${\hat{\mathbf{v}}}_i$转换到时域空间当中，以得到最终的音频分离结果：

$${\hat{\mathbf{s}}}_i={\mathcal{D}}_i\left({\hat{\mathbf{v}}}_i\right)=\text{ConvTr}1{\text{D}}_{C_{\mathcal{E}},K_{\mathcal{E}},K_{\mathcal{E}/2}}\left({\hat{\mathbf{v}}}_i\right)$$

参考链接

1. 【PyTorch】卷积层、池化层梯度计算 https://blog.csdn.net/weixin_44246009/article/details/119379516
2. 卷积神经网络-转置卷积 https://blog.csdn.net/weixin_38498942/article/details/106824520