An Efficient QP Variable Convolutional Neural Network Based In-loop Filter for Intra Coding

Inspired by channel attention module in [11], we propose a QPAM to avoid training and deploying multiple networks. Unlike channel attention module which extracts the attention map from the input feature map, our QP attention module is controlled by the QP value. The overview of the proposed QPAM is illustrated in Figure 1. Given a feature map $\mathbf{F}\in\mathbb{R}^{H\times W\times C}$ as input, QPAM sequentially infers a 1D QP attention map $\mathbf{M}\in\mathbb{R}^{1\times 1\times C}$. The attention process can be summarized as:

where $\otimes$ denotes element-wise multiplication. During multiplication, the attention values are broadcast along the channel dimension. $\mathbf{F^{\prime}}$ is the refined output. The process of generating the QP attention map is as follows: Given a QP value $q\in\Omega=[a,b]$, since $q$ is an integer, we first map $q$ to a vector $\mathbf{v}_{\Omega}(q)\in\mathbb{R}^{m\times 1}$ by one-hot encoding. $m$ is the length of $|\Omega|$. The QP attention map $\mathbf{M}$ is calculated by:

where $U\in\mathbb{R}^{C\times m}$ is a weight matrix and $\sigma(x)=log(1+e^{x})$. From the process we can see that the QPAM assigns different weights to each channel of the input feature map according to the QP value, so that the module can capture compression noise levels among different QPs. Meanwhile, the module can also emphasize meaningful features along channel axes through this process. Moreover, compared with other methods, the proposed QPAM has stronger scalability, which can also be easily applied to other discrete variables, e.g. frame types.

Architecture. In figure 2, we present the overview of the architecture of our QPALF network. This is also one of the popular architecture used by many other methods [8, 7, 9], which usually consists of three parts: the head part, the backbone part and the reconstruction part. The head part is responsible for initial feature extraction with only one convolutional layer followed by an activate function. Given a compressed input $X$, we can get a shallow feature $F_{0}$ through this layer:

The backbone part is the key component of the network, which is also the most distinct part of various networks. Here it makes up of $D$ cascaded residual feature aggregation modules (RFA). The backbone part receives the feature $F_{0}$ as input and sends the extracted global feature $F$ to the reconstruction part, which can be formulated as:

where $\mathcal{R}_{d}$ denotes the d-th RFA module function. $F_{d-1}$ is the input feature of the d-th RFA module function and $F_{d}$ is the corresponding output. The output features of the $D$ RFAs are concatenated together. Then we utilize a long skip connection to extract a global feature $F$. Finally, the global feature $F$ is transformed through the reconstruction part

where $\hat{Y}$ is the output and $\mathcal{H}$ is the reconstruction function, which consists of only one convolutional layer. A global residual learning is usually used to ease the training difficulty in the reconstruction part.

Inspired by [12], we propose a RFA module to make a better use of the local residual features. Figure 3 illustrates the detail of the RFA, which contains three residual blocks and one convolution layer. The input is extracted by three residual blocks at three different levels. Then the outputs from three residual blocks are concatenated and a $1\times 1$ convolution is applied at the end of the RFA for channel dimension reduction. As depicted in Figure 3, the detail of the residual block is the same as that in [13] except that we employ the proposed QPAM after the first convolution layer.

Considering the complexity of the proposed model, we use 6 RFA modules, that is, $D=6$ here. All of $3\times 3$ convolutional layers are followed by a parametric rectified linear unit (PReLU) [14] activation, and the filter number of each $3\times 3$ layer convolutional is 64. The $1\times 1$ convolutions are applied for channel dimension reduction.

Dataset. For the network training, we build a dataset using DIV2K [15] which contains 800 high-resolution images. First, we convert these images to YUV 4:2:0 color format and encode them by VTM6.0 [16] with all-intra (AI) configuration at four QPs, 22, 27, 32, 37. The built-in in-loop filters are all enabled when compressed these images. Then the compressed images are divided into two non-overlapping sets of training (700 images), validation (100 images). To further expand the training dataset, we split the reconstructions to small patches of $64\times 64$ with stride 16. And we remove the patches whose PSNR are more than 50.0 or less than 20.0. When training QP variable models, four training datasets are mixed in a random order.

Loss Function. Let $X$ be the input and $\theta$ be the set of network parameters to be optimized. Our goal is to learn an end-to-end function $F$ for generating a higher quality reconstruction $\hat{Y}=F(X;\theta)$ that is close to the ground truth $Y$. The loss function is the MSE between $\hat{Y}$ and $Y$:

where $N$ is the number of training samples in each batch. In order to train a more robust QP-combined network, we analyse the restoration ability of the network for different QPs. First we train a QPALF network using mixed dataset (the training detail will be presented in the follow). Then we plot the valid image cumulative proportion over the PSNR gain rate for different QPs. And the PSNR gain rate is defined as follows:

where $L_{init}$ is the MSE between $X$ and $Y$. From Figure 4, we can find: 1) The network has lower PSNR gain rate overall on dataset with smaller QP, especially at QP=22. Obviously smaller QP means less compression artifacts, and we do not expect the network to pay much attention on data with less compression artifacts. 2) The PSNR gain rate of 80% of the valid data is less than 10%, that is, the valid data with low PSNR gain rate accounts for a large proportion. So we expect the network to focus more on data that has low PSNR gain rate. To this end, we propose a focal MSE loss function to fine tune the proposed network, which can be calculated by:

where $\alpha_{q}$ is a weighting factor over QP value, and $\gamma$ is focusing parameter. Herein $\alpha_{q}=0.1,0.25,0.3,0.35$ for four QPs respectively and $\gamma=1$. Table 1 shows the coding performances of three networks over test sequences. QPALF-I, QPALF-II, QPALF-III denote QPALF without fine tuning, QPALF fine tuned by MSE and QPALF fine tuned by focal MSE respectively. As we can see, QPALF-III achieves more bit-rate saving than QPALF-II, which demonstrates the effectiveness of focal MSE.

Training Detail. The widely adopted deep learning framework Pytorch [17] is utilized to train our models. We use Adam [18] optimization to train these models and a batch size of 64. The learning rate discounts 0.5 every 25 epochs. The training takes 100 epochs in total. For QP-separate models, we first train the model for QP=37 and then use it to initialize the parameters of the networks with smaller QP. The initial learning rate is $10^{-4}$ for QP=37 and $10^{-5}$ for other QPs. For QP-combined models, the initial learning rate is $10^{-4}$ and the fine-tune process takes 50 epochs with learning rate $10^{-5}$. All models are trained on NVIDIA Titan X (Pascal) GPUs.

Implementation. We integrate the QPALF into VVC as an additional tool of in-loop filters between DBF and SAO. To get better performance for video coding, a frame level flag is signaled in the bitstream to indicate whether QPALF is enabled for this frame in the decoder. When the reduction of RD cost is greater than 0, the flag will be enabled and the QPALF will be applied for the frame on luma component.

In our experiments, all approaches for in-loop filtering are incorporated into the VVC reference software VTM6.0. The Libtorch [17] library is integrated into VTM 6.0 to perform the in-loop filtering with the different models. Four typical QP values are tested, including 22, 27, 32, 37. We use the AI configuration suggested by VVC common test condition (CTC) [19]. The anchor for all experiments is VTM6.0 with all built-in in-loop filters enabled. The coding efficiency is evaluated on standard video sequences from class A1 to class E recommended by JVET. The BD-Rate [20] are referred to measure the coding performance. We only train and apply the models on Y channel, but our approach can be extended to any arbitrary number of channels.

RD performance. First, we compare our QPALF with VTM baseline, and two CNN based in-loop filters, RHCNN [8] and DRNLF [9]. For a fair comparison, we also train the models using our dataset and integrate the trained model into VTM6.0 between DBF and SAO. The results are displayed in Table 2. It can be seen obviously, our QPALF further improves the coding efficiency, which obtains 4.03% bit-rate saving overall for the luma component on all the test sequences. To further verify the efficiency of the proposed model, we also compare our QPALF with QPMLF and QPALF-D. QPALF-D is the QPALF model trained separately on four QPs. QPMLF is the QPALF without QPAM and use the QP map method [10]. As we can see, compared with QP-separate model QPALF-S, the performance of the QP-combined model QPMLF degrades while our model QPALF achieves even better coding performance. Moreover, the PSNR gain of three models on multiple QPs is also depicted in Figure 5. We can observe that our QPAM obtains the highest PSNR gain over all QPs, which also demonstrates the generalization ability and robustness of the proposed method. (Since the models are trained on only four QPs, {22, 27, 32, 37}, we first map other QP values to the four QPs.)

Subjective evaluation. Figure 6 illustrates the subjective visual quality comparison among all four approaches. It can be observed that the images enhanced by our approach remain less distortion than those by other approaches, e.g., the clearer edge of the basketball net line. In Figure 7, we display the residual map of three methods over VTM baseline. Compared with RHCNN and DRNLF, our method can restore more image texture details.

Complexity. Table 3 shows the average encode/decode complexity increase and parameters of different models on Intel(R) Xeon(R) CPU E5-2697v4 and Titan X (Pascal). All of the neural networks are conducted with GPU acceleration. The complexity increase is calculated by $\Delta T=(\hat{T}-T)/T$, where $\hat{T}$ is the encode/decode time with integrating the models, $T$ is the original encode/decode time. Our proposed QPALF has less model parameters overall. Compared with QPMLF, our model achieves much better RD performance with little complexity increase.