Use AI to speed up AV1 encoder

发表于 | 阅读次数

0x1 AV1的编码复杂性

我们知道AV1的官方参考实现是libaom，由于AV1的编码复杂度高，如果采用libaom编码器来生成AV1的码流离实时编码还有很大的距离。我们知道传统视频编码器中有宏块的概念，宏块是16x16的亮度块 + 2个 8x8的色度快。从HEVC开始，到现在的AV1进一步引入了partiton的概念，也就是树形编码的概念，也就是说把先规定好最大的编码单元，这个最大的编码单元称为super block，在HEVC一般是64x64，在AV1中为128x128。然后进行四叉树划分，AV1中直至划分成4x4，HEVC中直至划分成8x8。而且这种划分进一步扩展到预测单元和变换单元。据统计，AV1编码中的复杂度80%是因为partiton引入的，所以要是能有一个快速方法来加速partiton的判断的话，AV1的编码速度能大幅提升。

AV1中的partiton划分如下图所示。

0x2 AI加速partiton划分

在libaom目前的实现中，AI加速主要用在intra frame的partition划分优化上。

1
2
3
4
av1_intra_mode_cnn_partition(
    &cpi->common, x, bsize, x->quad_tree_idx, &partition_none_allowed,
    &partition_horz_allowed, &partition_vert_allowed, &do_rectangular_split,
    &do_square_split);

该函数的输入是图像的像素值，可以理解为图像对应的纹理。并且需要把对应亮度/色度值转换成0~1之间的浮点数。如下代码所示，c为亮度或色度分量的index，这里c为0，为亮度分量。max_val为亮度/色度分量的最大值255(假设为8bit yuv)。

1
2
3
for (int i = 0; i < height; ++i)
    for (int j = 0; j < width; ++j)
      input[i * in_stride + j] = (float)dgd[c][i * stride + j] / max_val;

输出是这几个变量partition_none_allowed，partition_horz_allowed，partition_vert_allowed，do_rectangular_split，do_square_split。用来对后续的partition的划分进行优化，从变量的名称可以看到这些变量会对后续的partition划分进行限制，也就是减少了partition的数目。

推理采用的网络是CNN + DNN的结合。
CNN是5层网络结构，网络定义如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
static const CNN_CONFIG av1_intra_mode_cnn_partition_cnn_config = {
    NUM_CNN_LAYERS,  // num_layers
    0,               // is_residue
    0,               // ext_width
    0,               // ext_height
    0,               // strict_bounds
    {
        {
            CNN_LAYER_0_IN_CH,                                // in_channels
            CNN_LAYER_0_WIDTH,                                // filter_width
            CNN_LAYER_0_WIDTH,                                // filter_height
            CNN_LAYER_0_OUT_CH,                               // out_channels
            CNN_LAYER_0_HORZ_STRIDE,                          // skip_width
            CNN_LAYER_0_VERT_STRIDE,                          // skip_height
            0,                                                // maxpool
            av1_intra_mode_cnn_partition_cnn_layer_0_kernel,  // weights
            av1_intra_mode_cnn_partition_cnn_layer_0_bias,    // bias
            PADDING_VALID,                                    // pad
            RELU,                                             // activation
            0,                                                // deconvolve
            0,                                                // branch
            BRANCH_NO_COPY,                                   // branch_copy_type
            BRANCH_NOC,        // branch_combine_type
            NO_BRANCH_CONFIG,  // branch_config
            NO_BN_PARAMS,      // bn_params
            -1,                // output_num
        },
        {
            CNN_LAYER_1_IN_CH,                                // in_channels
            CNN_LAYER_1_WIDTH,                                // filter_width
            CNN_LAYER_1_WIDTH,                                // filter_height
            CNN_LAYER_1_OUT_CH,                               // out_channels
            CNN_LAYER_1_HORZ_STRIDE,                          // skip_width
            CNN_LAYER_1_VERT_STRIDE,                          // skip_height
            0,                                                // maxpool
            av1_intra_mode_cnn_partition_cnn_layer_1_kernel,  // weights
            av1_intra_mode_cnn_partition_cnn_layer_1_bias,    // bias
            PADDING_VALID,                                    // pad
            RELU,                                             // activation
            0,                                                // deconvolve
            0,                                                // branch
            BRANCH_NO_COPY,                                   // branch_copy_type
            BRANCH_NOC,        // branch_combine_type
            NO_BRANCH_CONFIG,  // branch_config
            NO_BN_PARAMS,      // bn_params
            3,                 // output_num
        },
        {
            CNN_LAYER_2_IN_CH,                                // in_channels
            CNN_LAYER_2_WIDTH,                                // filter_width
            CNN_LAYER_2_WIDTH,                                // filter_height
            CNN_LAYER_2_OUT_CH,                               // out_channels
            CNN_LAYER_2_HORZ_STRIDE,                          // skip_width
            CNN_LAYER_2_VERT_STRIDE,                          // skip_height
            0,                                                // maxpool
            av1_intra_mode_cnn_partition_cnn_layer_2_kernel,  // weights
            av1_intra_mode_cnn_partition_cnn_layer_2_bias,    // bias
            PADDING_VALID,                                    // pad
            RELU,                                             // activation
            0,                                                // deconvolve
            0,                                                // branch
            BRANCH_NO_COPY,                                   // branch_copy_type
            BRANCH_NOC,        // branch_combine_type
            NO_BRANCH_CONFIG,  // branch_config
            NO_BN_PARAMS,      // bn_params
            2,                 // output_num
        },
        {
            CNN_LAYER_3_IN_CH,                                // in_channels
            CNN_LAYER_3_WIDTH,                                // filter_width
            CNN_LAYER_3_WIDTH,                                // filter_height
            CNN_LAYER_3_OUT_CH,                               // out_channels
            CNN_LAYER_3_HORZ_STRIDE,                          // skip_width
            CNN_LAYER_3_VERT_STRIDE,                          // skip_height
            0,                                                // maxpool
            av1_intra_mode_cnn_partition_cnn_layer_3_kernel,  // weights
            av1_intra_mode_cnn_partition_cnn_layer_3_bias,    // bias
            PADDING_VALID,                                    // pad
            RELU,                                             // activation
            0,                                                // deconvolve
            0,                                                // branch
            BRANCH_NO_COPY,                                   // branch_copy_type
            BRANCH_NOC,        // branch_combine_type
            NO_BRANCH_CONFIG,  // branch_config
            NO_BN_PARAMS,      // bn_params
            1,                 // output_num
        },
        {
            CNN_LAYER_4_IN_CH,                                // in_channels
            CNN_LAYER_4_WIDTH,                                // filter_width
            CNN_LAYER_4_WIDTH,                                // filter_height
            CNN_LAYER_4_OUT_CH,                               // out_channels
            CNN_LAYER_4_HORZ_STRIDE,                          // skip_width
            CNN_LAYER_4_VERT_STRIDE,                          // skip_height
            0,                                                // maxpool
            av1_intra_mode_cnn_partition_cnn_layer_4_kernel,  // weights
            av1_intra_mode_cnn_partition_cnn_layer_4_bias,    // bias
            PADDING_VALID,                                    // pad
            RELU,                                             // activation
            0,                                                // deconvolve
            0,                                                // branch
            BRANCH_NO_COPY,                                   // branch_copy_type
            BRANCH_NOC,        // branch_combine_type
            NO_BRANCH_CONFIG,  // branch_config
            NO_BN_PARAMS,      // bn_params
            0,                 // output_num
        },
    },
};

DNN是两层网络结构，网络定义如下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
static const NN_CONFIG av1_intra_mode_cnn_partition_branch_0_dnn_config = {
    BRANCH_0_NUM_DNN_FEATURES,
    BRANCH_0_NUM_LOGITS,
    BRANCH_0_NUM_DNN_LAYERS,
    {
        BRANCH_0_NUM_DNN_LAYER_0_UNITS,
        BRANCH_0_NUM_DNN_LAYER_1_UNITS,
    },
    {
        av1_intra_mode_cnn_partition_branch_0_dnn_layer_0_kernel,
        av1_intra_mode_cnn_partition_branch_0_dnn_layer_1_kernel,
        av1_intra_mode_cnn_partition_branch_0_logits_kernel,
    },
    {
        av1_intra_mode_cnn_partition_branch_0_dnn_layer_0_bias,
        av1_intra_mode_cnn_partition_branch_0_dnn_layer_1_bias,
        av1_intra_mode_cnn_partition_branch_0_logits_bias,
    },
};

0x3 推理模型的训练

以上编码过程中采用的CNN + DNN的网络是经过训练的，在libaom的代码中并没有提供这个网络结构的训练过程代码。参考HEVC partition优化可以大体知道其训练过程。应该是设计好网络以后，通过大量样本数据来训练得到该推理网络的。