使用CNN预测HEVC的CU分割 (4) -- 使用不同的神经网络架构进行训练

在使用CNN预测HEVC的CU分割 (3) -- 构建神经网络模型中，已经介绍了两种最基本的用于预测CU分割的神经网络架构。架构的设计很重要，一个不好的神经网络可能导致模型很难训练，所以这篇文章我准备介绍一下我想到的其它不同的模型，用于尝试。

我已经使用了一种神经网络模型进行了训练和测试，并将代码放在了GitHub，里面介绍了我的模型架构，以及验证结果：

GitHub - wolverinn/HEVC-CU-depths-prediction-CNN: Using convolutional neural networks to predict the Coding Units (CUs) depths in HEVC intra-prediction mode, in order to reduce the time of the encoding process in HEVC.

按照上一篇文章的思路，模型大体上可以分为两类，一种是模仿HEVC预测CU分割的思路，分别对每一层进行分割或者不分割的判断，训练多个二分类器；另一种思路就是直接根据每个64x64的CTU输入预测对应的16个label的输出。这篇文章都是基于第二种思路。

最直接的方法就是上一篇中提到的直接预测16个label，但是这样当然会很难训练，甚至无法训练。所以我想到了以下的一些变换方式，减少每个输入需要预测的label数量。

1. 对64x64的CTU中的每一个16x16的块，预测一个label

看到有论文这样做的，是直接不以64x64作为输入，而是以更小的16x16作为输入。根据分割信息和图片的对应关系，可以知道16x16对应的分割信息就是只有一个label。不过这样做的缺点就是，仅仅根据16x16理论上是不足以判断出0、1、2、3这四层的分割信息的，最多只能判断到了16x16的这一层，也就是第2层，是否继续往下分割。由于没有更上层的信息，理论上没法判断是0还是1。

所以我的方案是输入一个64x64的整个CTU，先经过卷积层和池化层，变为一个16x16xn大小的矩阵之后，再和当前需要判断的16x16拼接起来，共同经过接下来的卷积、池化、全连接层。

2. 对每一个32x32的块，预测一个label

一个完整的64x64大小的CTU由四个32x32的块组成，每个32x32的块又包含四个label信息。但是这里的想法是只预测出一个label信息：0，1或者2，如果是0或1说明这个32x32的块本身就不需要分割了，如果是2，说明需要继续分割为16x16，而至于分割成16x16之后还是否需要继续分割到第3层，也就是8x8，那么对于这种情况还需要再训练另一个分类器。所以这种方案一共要训练两个不同的网络。第一个网络用于从32x32的输入预测一个label作为输出，可能取值为0，1，2.第二个网络用于在前一个网络输出为2的情况下，进一步判断分割出来的每个16x16的块是否需要分割，是一个二分类器。

3. 对每一个32x32的块，预测四个label

这个方案只需要训练一个网络，就能预测出0、1、2、3的结果作为输出。之前用64x64作为输入的时候，要预测16个label，而一个64x64可以分成四个32x32，因此一个32x32就对应4个label，每个label的取值是0、1、2、3. 这个方法其实就是在直接用64x64预测16个label的基础上缩小了规模，并且比起直接用16x16预测一个label这个方法还是有一定道理的，因为理论上32x32只能判断1、2、3这三个的情况，对于上层的64x64，其实是没法判断到底要不要作分割的。不过，这种方案下，只有结果中四个label全都预测为0，才能是0，稍微作了一些限制，所以还是可以先训练试试看一下最终的结果是否有效

以上的几种方案，都可以用ResNet再试一下