论文笔记 SuDORMRF：EFFICIENT NETWORKS FOR UNIVERSAL AUDIO SOURCE SEPARATION

SUDORMRF: EFFICIENT NETWORKS FOR UNIVERSAL AUDIO SOURCE SEPARATION

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

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

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

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

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

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

采用常规卷积：

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

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

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

2.Methodology

Overall Architecture

整体算法流程如下：

定义输入输出尺寸
x ∈ R T \mathbf{x}\in\mathbb{R}^T x∈RT为混合音频信号， E \mathcal{E} E为Encoder，对输入 x \mathbf{x} x处理得到特征向量: v x = E ( x ) ∈ R C ε × L \mathbf{v}{\mathbf{x}}=\mathcal{E}\left(\mathbf{x}\right)\in\mathbb{R}^{C{\mathbf{\varepsilon}}\times L} vx=E(x)∈RCε×L。将 v x \mathbf{v}_{\mathbf{x}} vx送入Separation Module S S S当中得到 m ^ i ∈ R C ε × L \hat{\mathbf{m}}i\in\mathbb{R}^{C{\mathcal{\varepsilon}}\times L} m^i∈RCε×L为第 i i i个音源的mask， i = 1 , 2 , ⋯ , N i = 1,2,\cdots,N i=1,2,⋯,N。假设共有 N N N个音源产生的音频信号 s 1 , s 2 , ⋯ , s N ∈ R T \mathbf{s}_1,\mathbf{s}_2,\cdots,\mathbf{s}N \in \mathbb{R}^T s1,s2,⋯,sN∈RT共同组成 x \mathbf{x} x。将 v x \mathbf{v}{\mathbf{x}} vx与 m i ^ \hat{\mathbf{m}_i} mi^逐项相乘得到第 i i i音源的特征向量 v i ^ \hat{\mathbf{v}_i} vi^：
v i ^ = v x ⊙ m i ^ \hat{\mathbf{v}i} = \mathbf{v}{\mathbf{x}} \odot \hat{\mathbf{m}_i} vi^=vx⊙mi^

再经过解码器 D \mathcal{D} D得到 s ^ i = D ( v i ^ ) \hat{\mathbf{s}}_i = \mathcal{D}(\hat{\mathbf{v}_i}) s^i=D(vi^)

定义1： Conv1D C , K , S : R C i n × L i n → R C × L \text{Conv1D}{C, K, S}: \mathbb{R} ^{C{in}\times L_{in}}\to \mathbb{R} ^{C\times L} Conv1DC,K,S:RCin×Lin→RC×L。表示一维常规卷积。将输入shape从 R C i n × L i n \mathbb{R} ^{C_{in}\times L_{in}} RCin×Lin转为 R C × L \mathbb{R} ^{C\times L} RC×L。其中 C C C为output channel， S S S为stride， K K K为kernel size，L为尺度（时间）。
定义2： ConvTr1D C , K , S : R C i n × L i n → R C × L \text{ConvTr1D}{C, K, S}: \mathbb{R} ^{C{in}\times L_{in}}\to \mathbb{R} ^{C\times L} ConvTr1DC,K,S:RCin×Lin→RC×L。转置卷积。需要注意的一点是，转置卷积的运算过程，相当于原卷积核对input求梯度（具体过程参看参考链接）。
定义3： DWConv1D C , K , S : R C i n × L i n → R C × L \text{DWConv1D}{C, K, S}: \mathbb{R} ^{C{in}\times L_{in}}\to \mathbb{R} ^{C\times L} DWConv1DC,K,S:RCin×Lin→RC×L。一维Depthwise Convolution深度卷积。

本质是将原始卷积层拆分成 G = C in G = C_\text{in} G=Cin 个Conv1D： F ^ i = [ Conv1D C G , K , S ] i \hat{\mathcal{F}} _i= [\text{Conv1D} _{C_G, K, S}] _i F^i=[Conv1DCG,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]。每一个卷积核 F ^ \hat{\mathcal{F}} F^对输出贡献 C G C_G CG个通道。最终将 { 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)
Concat ( ⋅ ) \text{Concat}(\cdot) Concat(⋅)表示拼接。

2.1 Encoder

Encoder用 E \mathcal{E} E表示，包含一个一维卷积，kernel size为 K E K_{\mathcal{E}} KE, stride为 K E / 2 K_{\mathcal{E}}/2 KE/2，使用公式表示Encoder的具体操作如下：
v x = E ( x ) = R e L U ( C o n v 1 D C E , K E , K E / 2 ( x ) ) ∈ R C E × L (3) \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} \tag3 vx=E(x)=ReLU(Conv1DCE,KE,KE/2(x))∈RCE×L(3)

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

2.2 分离模块

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

使用LN和Pointwise Conv将 v x \mathbf{v}x vx映射到新的通道空间当中：
y 0 = C o n v 1 D C , 1 , 1 ( L N ( v x ) ) ∈ R C × L (4) \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} \tag4 y0=Conv1DC,1,1(LN(vx))∈RC×L(4)

LN ( v x ) \text{LN}(\mathbf{v_x}) LN(vx)为layer-norm layer。
使用多个BU-convolutional blocks（U-ConvBlocks）拼接在一起。其中第 i i i个Block的输出作为第 i + 1 i+1 i+1个Block的输入。U-ConvBlock的具体细节参考Sec2.2.1，类似TDA-Net和U-Net，利用了多尺度信息，Block输入输出尺度一致。
使用最后的Block输出 y B T ∈ R L × C \mathbf{y}_B^T\in\mathbb{R}^{L\times C} yBT∈RL×C，针对每一个音源设置一个Conv1D层，以得到对应音源的中间特征向量 z i \mathbf{z}_i zi（比如有N个音源，则对应N个Conv1D，随后使用 z i \mathbf{z}_i zi获得 m i ^ \hat{\mathbf{m}i} mi^）：
z i = Conv 1 D C , C E , 1 ( y B T ) T ∈ R C E × L (5) \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} \tag5 zi=Conv1DC,CE,1(yBT)T∈RCE×L(5)

其中 y B T \mathbf{y}_B^T yBT表示 y B \mathbf{y}_B yB的转置。
使用 z i \mathbf{z}_i zi 求解 m i ^ \hat{\mathbf{m}i} mi^：
m ^ i = v e c − 1 ( exp ⁡ ( v e c ( z i ) ) ∑ j = 1 N exp ⁡ ( v e c ( z j ) ) ) ∈ R C E × L (6) \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} \tag6 m^i=vec−1(∑j=1Nexp(vec(zj))exp(vec(zi)))∈RCE×L(6)

其中 m ^ i ∈ [ 0 , 1 ] C E × L \hat{\mathbf{m}}{i} \in [0,1]^{C{\mathcal{E}}\times L} m^i∈[0,1]CE×L。 vec ( ⋅ ) : R K × N → R K ⋅ N \left ( \cdot \right ) : \mathbb{R} ^{K\times N}\to \mathbb{R} ^{K\cdot N} (⋅):RK×N→RK⋅N，表示向量化。 v e c − 1 ( ⋅ ) : R K ⋅ N → \mathrm{vec}^{- 1}\left ( \cdot \right ) : \mathbb{R} ^{K\cdot N}\to vec−1(⋅):RK⋅N→ R K × N \mathbb{R}^{K\times N} RK×N表示反向量化。可以发现，对于所有mask的同一位置，求和为1：
∑ i N m i ^ ( x , y ) = 1 \sum\limits_i^N {\hat{\mathbf{m_{i}}}_{ (x,y)}} = 1 i∑Nmi^(x,y)=1

其中 ( x , y ) (x,y) (x,y)表示mask的某个元素的坐标。

利用 m ^ i \hat{\mathbf{m}}_i m^i以及特征向量 v x \mathbf{v}_x vx得到每个音源的特征向量 v ^ i \hat{\mathbf{v}}_i v^i：
v ^ i = v x ⊙ m ^ i ∈ R C E × L (7) \hat{\mathbf{v}}i=\mathbf{v}\mathbf{x} \odot \hat{\mathbf{m}}i\in\mathbb{R}^{C{\mathcal{E}}\times L} \tag7 v^i=vx⊙m^i∈RCE×L(7)

2.2.1 U-ConvBlock

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

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

定义4： P R e L U C : R C × L → R C × L \mathrm{PReLU}_C:\mathbb{R}^{C\times L}\to\mathbb{R}^{C\times L} PReLUC:RC×L→RC×L。（parametric

rectified linear unit）：
P R e L U C ( y ) i , j = max ⁡ ( 0 , y i , j ) + a i ⋅ min ⁡ ( 0 , y i , j ) (8) \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) \tag8 PReLUC(y)i,j=max(0,yi,j)+ai⋅min(0,yi,j)(8)
a i \mathbf{a}_i ai为可学习的参数， y \mathbf{y} y为输入。
定义5： I M : R C × L → R C × M ⋅ L \mathcal{I}_M:\mathbb{R}^{C\times L}\to\mathbb{R}^{C\times M\cdot L} IM:RC×L→RC×M⋅L。上采样操作，最邻近插值， M M M为缩放系数。

2.3 Decoder

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

s ^ i = D i ( v ^ i ) = ConvTr 1 D C E , K E , K E / 2 ( 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) s^i=Di(v^i)=ConvTr1DCE,KE,KE/2(v^i)

参考链接

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