AV1基于机器学习的变换块快速划分

AV1基于机器学习的变换块快速划分

在上一篇文章《AV1基于机器学习的快速变换模式选择》中讲解了AV1如何使用机器学习技术为每个变换块选择合适的变换模式，本节将讲解AV1如何利用机器学习技术对变换块进行划分。

AV1无需像VP9中那样强制固定变换单元大小，而是允许编码块进行递归划分。并且支持从4×4到64×64的正方形，2:1/1:2和4:1/1:4比例也都可以。

如上图所示，一个编码块可以按多种方式划分为更小的TU（变换单元），每个小的TU还可以进一步划分。为了寻找最好的划分模式，编码器不得不为每个TU计算其划分和不话费的RD cost，从中选择RD cost最小的模式。

为了使用ML加速上述过程，可以将上面的问题转化为一个二元分类问题。目标就是构建一个ML模型，预测每个编码块是否进行TU划分。

ML模型的输入是残差块像素计算得到的一些特征。这些特征包括两个层级的均值和标准差。第一级是整个残差块的均值和标准差。第二级是每个1/4子块（对于正方形块和1：4矩形块）或每个1/2子块（对于1：2矩形块）的均值和标准差。同时还有计算第二级标准差的均值和均值的标准差。所以对于正方形块和1：4矩形块有12个特征，对于1：2矩形块有8个特征。（注：AV1代码中1：4矩形块其实使用了8个特征）为了预测更准确，对不同尺寸的块分别构建模型。所有模型的结构都如下图所示：

模型通过全连接神经网络实现，隐藏层使用ReLU作为激活函数。

模型使用方式如上图所示，对每个预测残差块提取特征，将特征输入模型获取输出。根据模型输出结果选择划分模型，如果模型认为很大概率不划分结果最优则不划分，如果模型认为很大概率划分结果最优则划分。其中预定义的阈值控制最终选择（AV1阈值默认为4000）。

具体模型结构

下面介绍两种具体的模型结构定义：

1、4x8块的模型

由前面知4x8模型输入8个特征，即输入单元数为8。它只有1个隐藏层，隐藏层单元数为16。模型参数如下：

//输入层到隐藏层间的权重weight
static const float av1_tx_split_nn_weights_4x8_layer0[8 * 16] = {
  0.068650f,  -0.732073f, -0.040361f, 0.322550f,  -0.021123f, 0.212518f,
  -0.350546f, 0.435987f,  -0.111756f, -0.401568f, 0.069548f,  -0.313000f,
  0.073918f,  -0.373805f, -0.775810f, -0.124753f, 0.181094f,  -0.602641f,
  -0.026219f, -0.350112f, 0.020599f,  -0.311752f, -0.476482f, -0.669465f,
  -0.310921f, 0.348869f,  -0.115984f, 0.154250f,  0.200485f,  -0.016689f,
  0.020392f,  0.413810f,  0.634064f,  -0.627530f, 0.399178f,  -0.012284f,
  0.472030f,  0.091087f,  -0.706100f, -0.447944f, -0.274226f, 0.445656f,
  0.309339f,  0.505522f,  0.038496f,  -0.152809f, 0.408684f,  -0.068151f,
  0.271612f,  0.353233f,  -0.150365f, 0.075212f,  -0.035096f, 0.346615f,
  0.124382f,  0.477072f,  0.216288f,  0.070548f,  -0.106362f, 0.681613f,
  -0.145502f, -0.218631f, -0.099248f, -0.001983f, -0.196819f, -0.969045f,
  0.063009f,  -0.123053f, 0.104875f,  -0.137581f, -0.282933f, -0.003624f,
  -0.315659f, -0.333523f, -0.503000f, -0.100063f, -0.536711f, -0.059978f,
  -0.670248f, -0.353762f, 0.181109f,  0.289715f,  -0.071206f, 0.261141f,
  0.052796f,  -0.114554f, -0.139214f, -0.261380f, 0.075984f,  -0.647925f,
  -0.099528f, -0.677814f, 0.015712f,  -0.389385f, -0.095622f, -0.165117f,
  -0.109454f, -0.175240f, -0.393914f, 0.212330f,  0.037822f,  0.248280f,
  0.180197f,  0.110493f,  -0.525727f, -0.092329f, -0.524029f, -0.407364f,
  -0.542373f, -0.435626f, -0.912194f, 0.062794f,  0.160433f,  0.741485f,
  -0.103659f, -0.119327f, -0.055275f, 0.334358f,  0.014713f,  0.046327f,
  0.831114f,  -0.576682f, 0.354369f,  -0.082088f, 0.452331f,  0.039730f,
  -0.792429f, -0.385862f,
};
//输入层到隐藏层间的偏置值bias
static const float av1_tx_split_nn_bias_4x8_layer0[16] = {
  0.238621f,  2.186830f,  1.383035f,  -0.867139f, 1.257119f, -0.351571f,
  -0.240650f, -0.971692f, 2.744843f,  1.116991f,  0.139062f, -0.165332f,
  0.262171f,  -1.598153f, -1.427340f, -1.602306f,
};
//隐藏层到输出层的权重
static const float av1_tx_split_nn_weights_4x8_layer1[16] = {
  -0.367134f, 1.373058f, -0.897039f, -0.326819f, -0.734030f, -0.290413f,
  -0.501249f, 0.505321f, -0.537692f, -0.767893f, 0.268697f,  0.278987f,
  0.085082f,  0.614986f, 0.847904f,  0.637578f,
};
//隐藏层到输出层的偏置
static const float av1_tx_split_nn_bias_4x8_layer1[1] = {
  0.20586078f,
};

2、8x8块的模型

由前面知4x8模型输入12个特征，即输入单元数为12。它只有1个隐藏层，隐藏层单元数为12。模型参数如下：

//输入层到隐藏层间的权重weight
static const float av1_tx_split_nn_weights_8x8_layer0[144] = {
  0.177983f,  -0.938386f, -0.074460f, -0.221843f, -0.073182f, -0.295155f,
  -0.098202f, -0.279510f, 0.001054f,  -0.119319f, -1.835282f, -0.581507f,
  -1.222222f, -1.049006f, -0.807508f, -0.454252f, -0.774879f, -0.180607f,
  -0.886976f, -0.231971f, -0.824677f, -0.351872f, -1.323819f, 0.235378f,
  0.015331f,  -0.341818f, 0.145549f,  -0.348362f, 0.147647f,  -0.323400f,
  0.047558f,  -0.553025f, -0.295485f, -0.330368f, -0.530605f, -0.407516f,
  0.447740f,  0.782381f,  -0.179164f, -0.584675f, -0.052645f, 0.038656f,
  -0.096783f, 0.038342f,  -0.170762f, -0.405844f, -0.552665f, -0.509866f,
  0.757204f,  -1.296465f, 0.631015f,  0.009265f,  0.646192f,  0.044523f,
  0.653161f,  0.033820f,  0.849639f,  -0.068555f, -1.036085f, -0.511652f,
  0.104693f,  -1.458690f, 0.286051f,  -0.089800f, 0.381564f,  -0.302640f,
  0.304465f,  -0.268706f, 0.432603f,  -0.117914f, -2.070031f, -0.565696f,
  -0.073027f, -1.783570f, -0.318144f, -0.320990f, -0.343966f, -0.140996f,
  -0.322977f, -0.232147f, -0.373210f, -0.158266f, -1.922305f, -0.634373f,
  0.101894f,  -0.221847f, 0.018412f,  -0.423887f, -0.266684f, -0.444930f,
  -0.196237f, 0.106638f,  -0.065834f, -0.538401f, -0.280772f, -0.620348f,
  1.089957f,  -0.799928f, 0.504112f,  -0.165763f, 0.578741f,  -0.172653f,
  0.547316f,  -0.143484f, 0.717220f,  -0.297190f, -1.237854f, -0.074819f,
  -0.977304f, -0.484092f, -0.646427f, -0.451443f, -0.612126f, -0.224475f,
  -0.731608f, -0.257077f, -0.665857f, -0.346742f, -1.216372f, 0.227267f,
  0.231249f,  -1.693073f, -0.035899f, 0.380845f,  -0.058476f, 0.409405f,
  -0.066679f, 0.406731f,  -0.068501f, 0.396748f,  0.639462f,  0.150834f,
  -0.418659f, -1.421931f, 0.101889f,  0.083573f,  0.129746f,  0.134460f,
  0.081185f,  0.127420f,  0.083664f,  0.051096f,  1.361688f,  0.386093f,
};
//输入层到隐藏层间的偏置值bias
static const float av1_tx_split_nn_bias_8x8_layer0[12] = {
  4.280443f, 2.218902f, -0.256953f, 3.161431f,  2.082548f, 2.506052f,
  2.563224f, 1.421976f, -1.627813f, -1.436085f, 2.297265f, 1.500469f,
};
//隐藏层到输出层的权重
static const float av1_tx_split_nn_weights_8x8_layer1[12] = {
  1.178833f,  -0.428527f, -0.078737f, 0.381434f, -0.466895f, -0.901745f,
  -0.766968f, -0.356663f, 0.450146f,  0.509370f, -0.356604f, -0.443506f,
};
//隐藏层到输出层的偏置
static const float av1_tx_split_nn_bias_8x8_layer1[1] = {
  -0.156294f,
};

训练集构建

在AV1中这些模型已经训练好了，训练模型的数据集也很好构建。选择不同类型、不同分辨率的序列，使用不同QP编码，在编码过程中计算输入特征（两级均值和标准差）及对应标签（RD cost）。

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