基于学习的编码（四）：WSE-CNNLF

基于学习的编码（四）：WSE-CNNLF

本文算法来自JVET-N0133，提出了一个CNN模型WSE-CNNLF(Wide-activated Squeeze-and-Excitation Convolutional Neural Network Loop Filter)来进行环路滤波。

与前面几篇算法不同的是，WSE-CNNLF并不是在VVC环路滤波中增加一个阶段或替换某个阶段，而是完全取代VVC的环路滤波（包括DBF、SAO、ALF）。

网络的输入和输出

WSE-CNNLF包括6个输入，其中三个是YUV重建图像，三个辅助输入。

（1）三个重建分量

YUV的三个重建图像首先根据比特深度归一化到[0,1]，然后输入到网络。归一化方法如下：

 

（2）三个辅助输入

第一个辅助输入是QP，不同的QP会导致不同水平失真。将QP作为输入可以只训练一个模型来应用于所有QP的情况。QP要归一化为QPmap:

 

另外两个辅助输入是亮度（Y）分量和色度（UV）分量的CU划分信息。因为块效应主要由块划分引起，输入CU的边界信息可以利用神经网络的注意力机制更高效的处理边界。CU的划分信息转换成CUmaps后归一化处理，然后输入网络。例如，对一帧的每个CU，边界位置填充2，其他位置填充1，如下图。归一化因子为2，可以得到两个CUmap，一个是Y-CUmap，另一个是UV-CUmap。

 

处理模块

Wide-activated convolution在超分辨和降噪任务中表现非常好。它由一个宽的3x3卷积跟上ReLU和一个窄的1x1卷积组成。

SE (Squeeze-and-Excitation)用于对每个卷积层加权。它能够利用不同通道之间的复杂关系，并为每个通道生成一个权重因子。

给定一个HxWxC的特征图，SE包含下面4个步骤：

（1）每个通道根据全局平均池化(Global Average Pooling ,GAP)得到一个值。

 

（2）一个全连接层加上一个ReLU函数增加了必要的非线性。其输出通道复杂度也降低一定比值r。

（3）第二个全连接层加上一个sigmoid函数给每个通道一个[0,1]范围内的平滑门限比率。

（4）最后，每个通道使用门限比率缩放。

 

上图值参数k表示第一步wide convolution中的wide ratio。参数r表示SE中的reduce ratio。

对一个基本处理单元，如果输入和输出的通道数一样则可以将输入和输出通过skip connection连接起来直接学习残差加快收敛，如上左图。

网络结构

WSE-CNNLF的处理包括三个阶段，其工作流程如下图所示。

 

（1）第1阶段

YUV三个分量分别通过wide-SE-blocks处理，每个与其对应的CUmap融合，融合方式是将CUmap逐像素与其对应通道相乘。由于U和V的尺寸只有Y的一半，所以需要上采用进行尺寸对齐。

（2）第2阶段

不同通道的特征图连接在一起，然后经过几个wide-SE-blocks处理。

（3）第3阶段

三个通道再次分别处理生成最终残差图像，然后与原始图像相加。

训练及结果

训练时网络配置如下：

 

 

上面是禁用DBF、SAO、ALF，仅开启WSE-CNNLF方法和VTM4.0比较的结果。

感兴趣的请关注微信公众号Video Coding