AV1基于机器学习的快速变换模式选择

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_4X4TX_8X8,    // TX_8X8TX_16X16,  // TX_16X16TX_32X32,  // TX_32X32TX_64X64,  // TX_64X64TX_4X4,    // TX_4X8TX_4X4,    // TX_8X4TX_8X8,    // TX_8X16TX_8X8,    // TX_16X8TX_16X16,  // TX_16X32TX_16X16,  // TX_32X16TX_32X32,  // TX_32X64TX_32X32,  // TX_64X32TX_4X4,    // TX_4X16TX_4X4,    // TX_16X4TX_8X8,    // TX_8X32TX_8X8,    // TX_32X8TX_16X16,  // TX_16X64TX_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 verticalADST_DCT,           // ADST in vertical, DCT in horizontalDCT_ADST,           // DCT in vertical, ADST in horizontalADST_ADST,          // ADST in both directionsFLIPADST_DCT,       // FLIPADST in vertical, DCT in horizontalDCT_FLIPADST,       // DCT in vertical, FLIPADST in horizontalFLIPADST_FLIPADST,  // FLIPADST in both directionsADST_FLIPADST,      // ADST in vertical, FLIPADST in horizontalFLIPADST_ADST,      // FLIPADST in vertical, ADST in horizontalIDTX,               // Identity in both directionsV_DCT,              // DCT in vertical, identity in horizontalH_DCT,              // Identity in vertical, DCT in horizontalV_ADST,             // ADST in vertical, identity in horizontalH_ADST,             // Identity in vertical, ADST in horizontalV_FLIPADST,         // FLIPADST in vertical, identity in horizontalH_FLIPADST,         // Identity in vertical, FLIPADST in horizontalTX_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 transformNULL,                            // 32x32 transformNULL,                            // 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 transformNULL,                            // 16x32 transformNULL,                            // 32x16 transformNULL,                            // 32x64 transformNULL,                            // 64x32 transform&av1_tx_type_nnconfig_4x16_hor,  // 4x16 transform&av1_tx_type_nnconfig_16x4_hor,  // 16x4 transformNULL,                            // 8x32 transformNULL,                            // 32x8 transformNULL,                            // 16x64 transformNULL,                            // 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 transformNULL,                            // 32x32 transformNULL,                            // 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 transformNULL,                            // 16x32 transformNULL,                            // 32x16 transformNULL,                            // 32x64 transformNULL,                            // 64x32 transform&av1_tx_type_nnconfig_4x16_ver,  // 4x16 transform&av1_tx_type_nnconfig_16x4_ver,  // 16x4 transformNULL,                            // 8x32 transformNULL,                            // 32x8 transformNULL,                            // 16x64 transformNULL,                            // 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