当我们训练好神经网络之后,下一步就是在HEVC的实际编码过程中使用这个神经网络模型提升编码效率,如果不能将神经网络集成到HEVC的编码器中进行编码,那我们训练这个编码的神经网络就失去了意义,并且,这样也不方便验证我们的模型到底能节省多少编码时间,以及BDBR和BD-PSNR的变化。
我已经将整个pipeline实现并且开源到了GitHub:
GitHub - wolverinn/HEVC-deep-learning-pipeline: Integrating neural network models in HEVC encoder, to test the complexity reduction using deep learning in HEVC intra-prediction.
要将模型部署到pipeline,首先要设计出pipeline的大致流程,并且,由于我们是在C++源码中嵌入Python运行神经网络模型,因此需要注意运行时间的问题,应该尽可能优化流程以提升编码效率,以下是我设计的流程:
首先,在HEVC编码器开始编码之前,加入一行代码启动Python脚本,这个Python脚本的功能是提取出要编码的YUV文件的每一帧,并且对每一帧使用神经网络模型预测对应的最佳分割深度。脚本以多线程方式启动,因为一个YUV文件有很多帧,预测不是一会儿就能做完的,如果每次编码都临时去预测或者等它全部预测完再启动编码,时间就会长很多。预测的结果存储在txt文件中,并且可以通过当前第几帧、第几个CTU找到。
在HEVC编码CTU时,通过索引找到对应的分割信息的预测值,当HEVC开始遍历所有深度时,如果当前所在深度和预测深度一致,则只计算当前深度,不继续向下分割,如果当前深度还未到达预测深度,则不计算当前深度,直接向下分割。
不计算当前深度,实现方式主要就是在xCompressCu
中,跳过xCheckRDCostIntra
这个函数。不继续分割的实现方式就是,阻止xCompressCu
的递归调用,直接把pcSubBestPartCU->getTotalCost()
赋值为最大的浮点数。
我在我的神经网络模型上测试了一下这个流程,编码时间大概能减少70%,所以这个设计还是很有效的。具体的改动可以去看我的GitHub,上面有更多细节。