2105_TIP_DeepQTMT：一种VVC帧内编码块划分的深度学习方法

论文题目：Deep QTMT_一种VVC帧内编码块划分的深度学习方法

来源：2021，TIP，BUAA

概述

需要解决的问题：简化在VVC中占据了编码时间的97%以上的编码块划分过程

本文的贡献：

1. 创建了大规模数据集
2. 提出了多阶段停止的CNN网络架构（MSE-CNN）以适应灵活的多阶段QTMT结构
3. 提出了自适应损失函数，综合了分割模式的不确定数量和最小化率失真损失的目标
4. 建立了一个多临界点的可调节编码复杂度和率失真损失的trade-off的决策框架

SOTA：相比于之前的SOTA方法，降低了VVC 44.65%-66.88%的编码时间，同时降低了bit-rate（BD-BR）1.322%-3.188%

Data Base: tianyili2017/CPIV: A large-scale database for QTMT-based CU partition of VVC

作者之前的工作：ETH-CNN，link of the repo

背景

VVC中的QTMT划分方式可分为QT（四叉树划分）、TT（不规则三叉树划分）和 BT（二叉树划分），如下图所示

其中，输入为一阵图片的亮度分量，不同阶段的CU使用不同的颜色进行区分

作者还进行了一个预实验，统计了不同尺寸的CU所对应的划分模式的比例：

从图中不难发现，不同尺寸的CU的划分模式比例差异较大，这说明我们应该设计不同类型的模块学习不同尺寸CU的特征，因此，这里引出一个问题：直接针对每个尺寸设计学习模块，会造成较大的参数冗余，如何将CU进行分类，设计模块对每个类别学习特征呢？

值得注意的是，这里因为VTM强制将128×128的CTU进行四叉树划分，因此所统计的最大尺寸的CTU就是64×64的。

作者的工作

创建了大规模数据集，链接地址见开头，论文中的TABLE1列出了数据集的配置情况。

网络架构

作者提出的网络的整体架构如下所示：

网络共分为6个Stage，分别对应CTU划分的六个阶段，对于每个阶段，首先CU会经过一个卷积层（Conditional Convolution），抽取纹理特征，然后输入到sub-net分支中，以预测划分模式，具体处理如下所示：

1. Conditional Convolution，这里的conditional主要是考虑到每层所要处理的CU大小不同，需要提取的特征也处于不同的level，所以结构上应该根据输入CU的尺寸进行动态调整——这就是预实验提出的问题。

   a. 根据当前编码块的粒度，从六种类型的残差块中选择一个，学习不同等级的特征

   Conditional Convolution中的残差块通过设置步长为1和zero-padding，保证了输入特征图和输出特征图的尺寸不变，因此这里我们只关注当前编码块（current CU）输入特征图的粒度（granularity）$a_c$，这里定义为$a=min(w,h)$。

   作者认为不同粒度的编码块需要抽取的纹理特征在不同的level上，设置了六个残差单元（如下图所示），编号1~6，分别对应粒度128~4的编码块，对应公式如下所示：
   $$k=lo{g}_2[\frac{256}{min(w,h)}]$$
   b. 根据parent CU和 current CU的划分方式决定残差块的数量

   Conditional Convolution中的残差块通过设置步长为1和zero-padding，保证了输入特征图和输出特征图的尺寸不变，因此设置几个残差块似乎并无大碍，但是作者这里仍然设置了一些规则，残差块的数量$n_r$由parent CU和current CU的粒度决定：
   $$n_r=\{\begin{array}{cc}lo{g}_2(\frac{a_p}{a_c}),& 4\le a_c\le 24\\ 1,& a_c=128\end{array}$$
   这里的$n_r$取整后在0，1，2之间取值。

   在实际操作中，我们将六个stage的Cond. Conv Layer设置为权值共享。

2. sub-net，这里的子网络用来预测划分模式，输入的CU会经过两层或三层的卷积层，以提取CU划分的low-level features——这里的卷积核的尺寸都是2的整数倍，例如2×2、4×4等，而且步长等于卷积核的尺寸，这里的non-overlapping的卷积操作对应于CU划分决策的non-overlapping特征；然后会被输入到全连接层中。

   按照输入CU粒度的不同，sub-net也有不同的类型，六个stage对应的sub-net如下所示

这里的卷积层的缩写是NC，即Non-overlapping convolution，下面以$N{C}_4$为例进行说明，如上图所示，该NC Layers共分为三层，分别对应的卷积核尺寸为4×4、2×2和2×2，每层的卷积核的感受野如下图所示，分别是4×4、8×8和16×16

作者也提供了一个动画，展示这个过程：https://www.bilibili.com/video/BV1Hq4y1F7KJ?share_source=copy_web

在non-overlapping convolution layer之后，就是全连接层，输出一个one-hot向量，这里输出的向量的长度在2-6之间，取决于CU的尺寸。

值得注意的是，QP对于CU的划分也有影响，当量化参数下降时，CU更加倾向于继续划分，因此这里将QP作为额外的参数添加到全连接层中，这里添加的方法不同于他们之前的直接将QP拼接到全连接层中，而是将一半的特征值乘以标准化之后的特征（half-mask），这里标准化QP的计算公式如下
$$\hat{q}=\frac{q}{51}+0.5$$

Loss Function

这里我们考虑三个问题：

1. 在stage2-6中，不同尺寸的编码单元有着不同的候选划分模式，即某些划分模式对于某些尺寸的CU是非法的；
2. 在预实验中，我们看到了不同划分模式之间的分布比例差异悬殊；
3. 在VVC中，不同的划分模式可能对应不同的RD Cost，简单的交叉熵并不能体现这一点。

针对问题1，我们定义了对于宽为w，高为h的编码单元的合法划分模式集合为$M(w,h)$，其中 m ∈ { 0 , 1 , 2 , 3 , 4 , 5 } m\in\{0,1,2,3,4,5\} m∈{0,1,2,3,4,5}，（具体需要自己去写一套规则，文中并未说明），而对于一个mini-batch的CUs尺寸相同，因而合法划分模式集合也是相同的，于是损失函数可以被表示为
$$L_{CE,B}=-\frac{1}{N}\sum _{n=1}^N\sum _{m\in M}{y}_{n,m}log({\hat{y}}_{n,m})$$
这里的$y_{n,m}$和${\hat{y}}_{n,m}$分别代表第n个CU在划分模式m下的GT二进制标签和预测概率，下标B表示Basic。

针对问题2，我们设置了惩罚因子$a$控制是否考虑占比较小的划分模式，此时的损失函数为
$$L_{CE}=-\frac{{\sum }_{n=1}^N(\frac{1}{p_m}{)}^a{\sum }_{m\in M}{y}_{n,m}log({\hat{y}}_{n,m})}{{\sum }_{n=1}^N(\frac{1}{p_m}{)}^a}$$
当a=0时，完全不考虑划分模式分布差异巨大的事实，当a=1时，占比越小的项在训练中的重要性越大，作者推荐a=0.3。

针对问题3，上面的式子仅仅让预测模式尽量贴合最佳划分模式，但是RD Cost却并未控制，于是作者提出增加$L_{RD}$项，尽量让率失真损失也越小
$$L_{RD}=-\frac{1}{N}\sum _{n=1}^N\sum _{m\in M}{y}_{n,m}(\frac{r_{n,m}}{r_{n,min}}-1)$$
上式括号中的项可视为率失真损失的normal value，然后综合的损失函数为
$$L=L_{CE}+\beta \cdot L_{RD}$$

Multi-Threshold Decision Mechanism

理想的情况是MSE-CNN能够取得和暴力搜索相同的结果，但这是不可能达到的，即因为模型自身限制导致的 RD degradation 是不可避免的，例如最佳和次佳模式的评分十分接近，但是快速决策只能选择其中一种模式，为此我们提出了multi-threshold 框架，即可以根据用户的需求在编码复杂性和RD performance 之间取得平衡。

这里的threshold表示为${\tau }_{s=2}^6$， 取值在[0,1]，分别代表Stage2-Stage6的决策threshold，而作者所谓的multi体现在不同阶段的threshold各不相同。

具体操作为，假设${\hat{y}}_{n,m}$表示第n层的CU被划分为模式m的概率，而${\hat{y}}_{n,max}$表示第n层的CU最有可能被划分模式下的概率值，即 y ^ n , m a x = m a x m ∈ M { y ^ n , m } \hat{y}_{n,max}=max_{m\in M}\{\hat{y}_{n,m}\} y^​n,max​=maxm∈M​{y^​n,m​}，只有当sub-net输出的概率值${\hat{y}}_{n,m}$大于${\tau }_s\cdot {\hat{y}}_{n,max}$的时候，才对其对应的模式m进行RDO判断孰优孰劣。

这样，就可以在快速决策和RDO之间取得一个trade off，例如当τ=1的时候，就是完全的快速决策，当τ=0的时候，就是完全的RDO。

作者通过在CPIV 数据集上的测试，得到了不同阶段的预测准确性和threshold的关系，如下图所示

可见，stage2和stage6的预测准确性随着threshold的提高而相对较慢，其他阶段的准确性下降较快，说明前两者可以较多的偏向于快速决策，作者并给出了设置threshold的策略：

1. 一种节省时间的策略，当平均threshold>=0.4时，$\tau 2>\tau 6>\tau 3=\tau 4=\tau 5$
2. 一种保证RD Performance的策略，当平均threshold<0.4时，$\tau 2>\tau 4=\tau 3=\tau 5=\tau 6$

实验

实验参数设置

所有实验均在VTM7.0上进行，测试集是CPIV的800张测试图片和JVET上classes A~E的22个视频序列，模式均为全帧内（encoder_intra_vtm.cfg），量化参数为{22，27，32，37}。

作者分别在luma和chroma两个通道训练MSE-CNN，所有超参数都在验证集文件中，使用Xavier initialization初始化所有权重和偏置，学习率被初始化为10^-4 , 并每2000次迭代时指数化下降1%，一共训练50万次。

训练步骤如下，每个模块被训练1~2h，总共训练时长为30h

作者还按照faster、fast和medium设置了不同类型的multi-threshold，平衡效果和复杂度。

作者还将以往的SOTA方法的复杂度和RD performance的表现和所提出方法进行了对比，并进行了消融实验验证了多阶段提前终止、损失函数中的率失真损失项和multi-threshold的合理性。

总结

本文的核心思想就是通过引入“粒度”的概念，通过预实验发现不同粒度的划分模式比例存在显著差异，并在网络设计中，针对不同粒度的编码单元的特征提取模块设计了不同类型的Conditional Convolution和Non-overlapping Convolution，针对因为编码块划分造成的粒度变化设计了Conditional Convolution Layer中的残差块数量。

此外，本文还充分考虑了多阶段（即提前终止机制）和 Multi-threshold机制（尽量减轻RD Degradation，并通过实验结果给出了两种策略）。

对我们的借鉴意义包括：

1. 多阶段输出，即提前终止机制可以通过sub-net的方式实现，这里通过half-mask的方式将QP的影响考虑进来，是一个比较有新意创新（之前作者是将QP作为超参数添加到全连接层的末尾），其余部分也可以借鉴
2. 借鉴作者创建数据集的方法（后面会写一篇博客总结下这个系列的数据集创建步骤和注意事项）
3. 思考损失函数的论述很精彩，作者关注到了不同划分模式下的标签分类差异较大的实验事实，并添加了RD Cost的项，之后通过消融实验证实了这个trick确实有效