AV1基于机器学习的快速变换模式选择
1）变换块分区：AV1无需像VP9中那样强制固定变换单元大小，而是允许亮度间编码块划分为多种大小的变换单元，这些递归分区最多可递减2级。为了合并AV的扩展编码块分区，我们支持从4×4到64×64的正方形，2:1/1:2和4:1/1:4比例也都可以。此外，色度转换单元总是要尽可能地大。
所以支持的变化尺寸如下：
static const TX_SIZE txsize_sqr_map[TX_SIZES_ALL] = {
TX_4X4, // TX_4X4
TX_8X8, // TX_8X8
TX_16X16, // TX_16X16
TX_32X32, // TX_32X32
TX_64X64, // TX_64X64
TX_4X4, // TX_4X8
TX_4X4, // TX_8X4
TX_8X8, // TX_8X16
TX_8X8, // TX_16X8
TX_16X16, // TX_16X32
TX_16X16, // TX_32X16
TX_32X32, // TX_32X64
TX_32X32, // TX_64X32
TX_4X4, // TX_4X16
TX_4X4, // TX_16X4
TX_8X8, // TX_8X32
TX_8X8, // TX_32X8
TX_16X16, // TX_16X64
TX_16X16, // TX_64X16
};其中矩形块的变换是通过更小的正方形变换实现的。
2）扩展的转换内核：为AV1中的帧内和帧间块定义了一组更丰富的转换内核。完整的2-D内核集由DCT，ADST，flipADST和IDTX 的16个水平/垂直组合组成。除了已在VP9中使用的DCT和ADST之外，flipADST则以相反的顺序应用ADST，并且恒等变换（IDTX）意味着沿某个方向跳过变换编码，因此对于编码锐利边缘特别有用。随着块大小变大，某些内核开始发挥类似作用，因此，随着变换大小的增加，内核集会逐渐减少。
DCT，ADST，flipADST和IDTX 的4个水平变换4个垂直变换可以组合成16个2D变换。其在libaom中定义如下：
enum {
DCT_DCT, // DCT in both horizontal and vertical
ADST_DCT, // ADST in vertical, DCT in horizontal
DCT_ADST, // DCT in vertical, ADST in horizontal
ADST_ADST, // ADST in both directions
FLIPADST_DCT, // FLIPADST in vertical, DCT in horizontal
DCT_FLIPADST, // DCT in vertical, FLIPADST in horizontal
FLIPADST_FLIPADST, // FLIPADST in both directions
ADST_FLIPADST, // ADST in vertical, FLIPADST in horizontal
FLIPADST_ADST, // FLIPADST in vertical, ADST in horizontal
IDTX, // Identity in both directions
V_DCT, // DCT in vertical, identity in horizontal
H_DCT, // Identity in vertical, DCT in horizontal
V_ADST, // ADST in vertical, identity in horizontal
H_ADST, // Identity in vertical, ADST in horizontal
V_FLIPADST, // FLIPADST in vertical, identity in horizontal
H_FLIPADST, // Identity in vertical, FLIPADST in horizontal
TX_TYPES,
DCT_ADST_TX_MASK = 0x000F, // Either DCT or ADST in each direction
} UENUM1BYTE(TX_TYPE);由上面可知当AV1为一个变换块选择划分方式和变换模式时需要遍历19x16=304种模式组合，复杂度非常高。
变换模式快速剪枝
为了减少变换模式决策的复杂度，AV1实现了基于机器学习（全连接神经网络FCN）的变换模式剪枝方法。在选择水平变换模式和垂直变换模式的时可以分别通过两个机器学习模型根据得分剪枝掉概率低的变换。
该机器学习模型是通过全连接神经网络FCN构建的，每层的模型参数（包括weights和bias）都预先训练好了直接在源码中固定了。AV1为水平变换和垂直变换分别都预先提供了19种机器学习模型，分别用于处理19种变化尺寸。
模型定义如下：
// Map tx_size to its corresponding neural net model for tx type prediction.
static NN_CONFIG_V2 *av1_tx_type_nnconfig_map_hor[] = {
&av1_tx_type_nnconfig_4x4_hor, // 4x4 transform
&av1_tx_type_nnconfig_8x8_hor, // 8x8 transform
&av1_tx_type_nnconfig_16x16, // 16x16 transform
NULL, // 32x32 transform
NULL, // 64x64 transform
&av1_tx_type_nnconfig_4x8_hor, // 4x8 transform
&av1_tx_type_nnconfig_8x4_hor, // 8x4 transform
&av1_tx_type_nnconfig_8x16_hor, // 8x16 transform
&av1_tx_type_nnconfig_16x8_hor, // 16x8 transform
NULL, // 16x32 transform
NULL, // 32x16 transform
NULL, // 32x64 transform
NULL, // 64x32 transform
&av1_tx_type_nnconfig_4x16_hor, // 4x16 transform
&av1_tx_type_nnconfig_16x4_hor, // 16x4 transform
NULL, // 8x32 transform
NULL, // 32x8 transform
NULL, // 16x64 transform
NULL, // 64x16 transform
};
static NN_CONFIG_V2 *av1_tx_type_nnconfig_map_ver[] = {
&av1_tx_type_nnconfig_4x4_ver, // 4x4 transform
&av1_tx_type_nnconfig_8x8_ver, // 8x8 transform
&av1_tx_type_nnconfig_16x16, // 16x16 transform
NULL, // 32x32 transform
NULL, // 64x64 transform
&av1_tx_type_nnconfig_4x8_ver, // 4x8 transform
&av1_tx_type_nnconfig_8x4_ver, // 8x4 transform
&av1_tx_type_nnconfig_8x16_ver, // 8x16 transform
&av1_tx_type_nnconfig_16x8_ver, // 16x8 transform
NULL, // 16x32 transform
NULL, // 32x16 transform
NULL, // 32x64 transform
NULL, // 64x32 transform
&av1_tx_type_nnconfig_4x16_ver, // 4x16 transform
&av1_tx_type_nnconfig_16x4_ver, // 16x4 transform
NULL, // 8x32 transform
NULL, // 32x8 transform
NULL, // 16x64 transform
NULL, // 64x16 transform
};下面以4x4的水平模型为例讲解其结构：
上图即为该模型的网络结构，输入输出各4个单元，包含一个隐藏层（8个隐藏单元）。有sigmoid和relu两种激活方式，默认不激活直接输出。
输入层到隐藏层的weights和bias如下：
static const float av1_tx_type_nn_weights_4x4_hor_layer0[32] = {
-1.64947f, -1.54497f, -1.62832f, -0.17774f, -2.89498f, -0.72498f, 0.72036f,
0.17996f, 1.20000f, -0.27654f, 0.77396f, 1.21684f, -1.75909f, -0.51272f,
-1.25923f, 0.35005f, -0.04257f, -0.23389f, -0.41841f, -0.08229f, 0.09503f,
2.73144f, -0.16875f, -0.23482f, 0.02194f, -0.26427f, 0.28049f, 0.21260f,
1.35792f, 0.27733f, 0.88660f, -0.68304f,
};
static const float av1_tx_type_nn_bias_4x4_hor_layer0[8] = {
1.38742f, 0.59540f, -1.37622f, 1.92114f,
0.00000f, -0.38998f, -0.32726f, -0.15650f,
};隐藏层到输出层的weights和bias如下：
static const float av1_tx_type_nn_weights_4x4_hor_layer1[32] = {
1.65254f, 1.00915f, -0.89318f, -2.05142f, -0.23235f, 0.96781f, -0.37145f,
-0.21056f, 1.13891f, 0.38675f, 0.87739f, -1.42697f, 0.48015f, 0.61883f,
-0.03979f, 0.11487f, 0.48042f, 0.45200f, -0.23242f, 0.75166f, 0.55458f,
0.39452f, -0.35285f, 1.59120f, -1.49221f, -0.48349f, -0.64692f, 1.49297f,
-0.26782f, -0.65416f, -0.10648f, 0.05568f,
};
static const float av1_tx_type_nn_bias_4x4_hor_layer1[4] = {
4.07177f,
3.26961f,
0.58083f,
1.21199f,
};变换块尺寸快速剪枝
AV1为每种尺寸块定义的机器学习模型如下：
static const NN_CONFIG *av1_tx_split_nnconfig_map[TX_SIZES_ALL] = {
NULL, // TX_4X4,
&av1_tx_split_nnconfig_8x8, // TX_8X8,
&av1_tx_split_nnconfig_16x16, // TX_16X16,
&av1_tx_split_nnconfig_32x32, // TX_32X32,
&av1_tx_split_nnconfig_64x64, // TX_64X64,
&av1_tx_split_nnconfig_4x8, // TX_4X8,
&av1_tx_split_nnconfig_4x8, // TX_8X4,
&av1_tx_split_nnconfig_8x16, // TX_8X16,
&av1_tx_split_nnconfig_8x16, // TX_16X8,
&av1_tx_split_nnconfig_16x32, // TX_16X32,
&av1_tx_split_nnconfig_16x32, // TX_32X16,
&av1_tx_split_nnconfig_32x64, // TX_32X64,
&av1_tx_split_nnconfig_32x64, // TX_64X32,
&av1_tx_split_nnconfig_4x16, // TX_4X16,
&av1_tx_split_nnconfig_4x16, // TX_16X4,
&av1_tx_split_nnconfig_8x32, // TX_8X32,
&av1_tx_split_nnconfig_8x32, // TX_32X8,
&av1_tx_split_nnconfig_16x64, // TX_16X64,
&av1_tx_split_nnconfig_16x64, // TX_64X16,
};每个模型都是全连接神经网络FCN，其原理和前面的类似。通过这个FCN可以快速判断该块是否还需要进一步向下划分，节省了搜索每种划分方式的时间，可以降低复杂度。
总结
以上便是AV1通过机器学习进行变换模式快速选择的方法，其中每个FCN的层数和神经元数量都比较小所以进行神经网络计算的复杂度比较低。通过FCN进行剪枝可以免去对很多模式的遍历计算可以节省大量计算时间。这些FCN都是预先定义和训练好的，不需要通过码流传输。
感兴趣的请关注微信公众号Video Coding
