今天介绍一篇有关去马赛克的工作，去马赛克是 ISP 流程里面非常重要的一个模块，可以说是将多姿多彩的大千世界进行色彩还原的重要一步。这篇工作探索的是如何从各种各样的去马赛克算法中，选择最佳的一种。
Abstract
本文提出了一种方法，能在给定损失函数和训练数据的情况下，跨越一系列计算预算，自动合成高效、高质量的去马赛克算法。该方法执行多目标、离散连续优化，可同时求解出能在计算成本和图像质量之间实现最佳权衡的程序结构及参数。我们对该方法进行了设计，旨在利用特定领域的结构来提高搜索效率。
我们将其应用于多项任务，包括对拜耳（Bayer）和富士 X-Trans 彩色滤光片图案进行去马赛克处理，以及联合去马赛克和超分辨率处理。在 8 个图形处理器（GPU）上运行几天时间后，它生成了一系列算法，相较于先前的最先进技术，在从每像素 10 次到 1000 次运算的一系列计算预算范围内，显著提高了图像质量（在相同成本下，质量提高 1 分贝至 3 分贝；在相同或更好质量下，吞吐量提高 8.5 至 200 倍）。由此生成的程序将经典去马赛克算法和基于深度学习的去马赛克算法的特性结合起来，形成了更高效的混合组合，这些组合从构建上就具备带宽高效性且可矢量化。最后，我们的方法会自动对所有生成的程序进行调度并将其编译为适用于现代处理器的优化单指令多数据（SIMD）代码。
Introduction
- 图 1：我们采用学习与程序搜索相结合的方式，自动合成了高效、高质量的去马赛克算法。与先前的最先进方法相比，它们在每像素 10 次到 1000 次运算的范围内，显著推进了成本与质量的帕累托前沿（右图所示）。相对于先前已发布的算法，在相同成本下，我们的算法质量至少高出 1dB；在相同质量下，速度要快 5 - 10 倍。在具有挑战性的图像内容上，视觉质量有了显著提升（注意在线性最小均方误差（LMMSE）、梯度 Halide 以及可变数目梯度（VNG）输出中存在的拜耳网格斑点和拉链状伪影）。唯一比我们算法质量更高的先前方法是那些大型卷积模型，但它们的计算成本要高出 2 - 3 个数量级（如去马赛克网络（Demosaicnet）、亨茨等人的相关算法）。除了此处展示的传统拜耳去马赛克算法之外，我们还提出了针对从 X - Trans 传感器进行去马赛克以及针对联合去马赛克超分辨率和单独超分辨率的具有帕累托优势的算法。
去马赛克处理是最普遍且对性能要求极高的图像处理任务之一。作为关键的第一步，它能够决定整个相机成像流程的结果好坏：在去马赛克处理过程中丢失的任何细节都将永远消失，而引入的任何虚假细节都会变成一种几乎不可能在后续流程中去除的复杂空间结构。要平衡这两方面的问题很困难，而且这个问题是不适定的，所以不存在正确答案。
同时，去马赛克处理往往必须在极为有限的计算资源条件下进行：一条 4K、60 帧每秒的视频流每秒需要处理 5 亿像素。即便我们将高端移动处理器（中央处理器、图形处理器或数字信号处理器）中的一个核心专门用于去马赛克处理任务，并且实现了完美的单指令多数据（SIMD）利用率，这也最多只能为每个像素留下执行几百次运算的时间。
面对这一挑战，当前的去马赛克算法通常瞄准两个极端之一（图 1）。从手机到 Adobe Camera Raw 等应用中，大多数广泛部署的实现方式，都被限制在每个像素最多进行数百次高度优化的运算，这些运算是手工定制的，用于对单一特定的彩色滤光片阵列进行反演。通过这种方式，它们能提供尚可的图像质量，但在具有挑战性的情况下，很难避免诸如摩尔纹和虚假细节等伪影问题。
与此同时，基于深度学习和优化的方法已经出现，它们极大地提升了图像质量，并且更容易推广到不同的彩色滤光片阵列以及其他问题变体上，但代价是计算量要多出 2 到 3 个数量级（每个像素需要进行数十万到数百万次运算），这使得它们超出了大多数实际应用场景的可承受范围。根据所选的实现方式不同，去马赛克处理在 Adobe Camera Raw 图像信号处理器（ISP）运行时间中所占的比重可能在 25% 到 85% 之间。
我们开发了一系列新的高效的、基于学习的去马赛克算法，这些算法显著提升了在每个像素运算次数从数十次到数千次这一整个范围内所能达到的最先进的图像质量。除了常见的拜耳模式之外，我们还针对富士 X-Trans 模式开发了去马赛克算法，并且联合解决了基于拜耳模式的去马赛克和超分辨率问题。
我们的程序具有帕累托优势：与同运算量范围的任何现有算法相比，它们能在相同的计算成本下提供显著更高的质量（1 分贝至 3 分贝），并且能以大幅降低的计算成本（8.5 倍至 220 倍甚至更多）呈现出相当或更好的图像质量。它们专为高效的流式单指令多数据（SIMD）实现而设计，并且能自动编译为适用于现代处理器的高度优化的内核。
我们通过开发一种多目标、离散连续搜索自动生成了这一系列新算法，该搜索同时求解程序结构和参数，以便在目标计算预算范围内找到计算成本与图像质量之间的最佳权衡。这种搜索由与近期的去马赛克及超分辨率神经网络相同的损失函数和训练数据驱动。然而，我们发现标准的神经架构搜索（NAS）技术不足以完成我们的任务：这些方法通常针对的是高度规则且参数过多的模型。我们专注于低成本模型，这需要精心设计以利用特定领域的结构。我们的搜索在 8 块 GeForce Titan Xp 图形处理器上经过 4 到 5 天就能生成最先进的结果 ------ 这与训练单个神经网络直至收敛所需的成本处于同一量级。
最终生成的程序将经典的以及基于深度学习的去马赛克和超分辨率算法的特点结合起来，形成了更高效的混合组合，将构建模块组合成从构造上就具备带宽高效且高度可向量化特点的算法。最后，我们的方法会自动对搜索生成的任何程序进行调度并将其编译为高度优化的单指令多数据（SIMD）代码。
我们认为，我们的方法为自动优化图像处理流程的性能和质量奠定了基础，它结合了经典算法和深度学习的优势，能够生成比现有算法更好、更高效的算法。例如，除了去马赛克问题的三种变体之外，我们还展示了我们的搜索方法也能够针对高性能超分辨率这一任务生成具有帕累托优势的程序。
总计起来，我们做出了以下贡献：
- 提出了全新的、处于最前沿水平的拜耳和 X-Trans 去马赛克算法、结合超分辨率的联合去马赛克算法，以及独立的超分辨率算法，这些算法在与商业最相关的计算预算范围内，性能远远优于此前的相关成果。
- 提供了一种能够自动生成涵盖广泛计算预算范围的此类算法的方法。
- 我们证明了添加特定领域的基元和搜索结构，可显著改善在低成本情况下通过可微程序搜索处理图像处理任务时所能实现的性能与质量的权衡。
- 我们定义了一个搜索空间，该空间通过构造能生成便于单指令多数据（SIMD）处理且利于局部性优化的算法，还定义了一个编译器，该编译器利用这种结构自动生成高度优化的流式实现方式。
Method
- 图 2 系统概述：我们针对去马赛克程序定义了一个搜索空间，这些去马赛克程序由传统神经网络以及特定领域图像处理构建模块的有向无环图（DAG）构成。我们进一步将这个空间分解为针对绿色预测和色度预测子程序的搜索。我们使用遗传搜索与帕累托抽样相结合的方式来选择哪些程序要进行变异以及在各代之间进行保留。我们程序的帕累托优势是通过我们的成本模型（用于估算计算成本）以及一个低成本的训练过程（用于估算给定程序的质量）来衡量的。
为了生成能在成本与质量之间取得平衡的去马赛克算法，我们必须高效地对大量候选程序进行搜索。每个程序都是可微的且包含大量参数，并且需要在大型数据集上进行训练。因此，驱动我们做出以下设计决策的关键限制因素就是运行搜索所需的时间。我们采取了一些措施来减少训练每个程序所需的时间，并降低搜索空间的组合复杂性，使其朝着更有可能快速且高质量的程序方向发展。
我们使用一种遗传搜索算法，该算法用上一代中性能最佳的程序的变异体来填充每一代。图 2 展示了整个系统。由于我们要针对多个目标进行优化，所以 "性能最佳" 的概念并不简单。我们希望得到能覆盖一系列运行时间预算的高质量程序。因此，我们将成本轴划分为多个成本层级，然后根据各程序与帕累托前沿的接近程度，对每个层级的程序进行变异并加以维护。
我们通过在一个小数据集上采用快速训练流程来预估程序质量，并利用一个简单模型（所执行操作的加权总和）来预估计算成本。从最后一代中，我们选取按帕累托优势排名前 100 的程序，将它们编译为高效的 Halide（拉根 - 凯利等人 2012 年提出）实现方式以测量真实成本，并在整个数据集上对它们进行训练以获取真实质量。
Building Blocks
- 图 3 我们利用受经典边缘自适应（a）和基于深度学习（b）的去马赛克算法启发的构建模块来构建程序搜索空间。我们遵循经典模式，首先重建绿色（G），然后通过预测红色和蓝色（R/B）与重建后的绿色之间的差值来重建红、蓝颜色。经典的边缘自适应去马赛克算法通常会在针对绿色的方向性一维滤波器之间进行选择以适应边缘。因此，我们纳入了一维和二维分组卷积、SoftMax 层、逐元素乘法以及求和归约操作来重现这一过程。对于深度学习模式，我们引入了常用的打包和解包原语（b），这些原语可将拜耳模式打包为平移不变图像，并将它们解包回输入分辨率。我们还包含了标准的深度学习原语，包括卷积层和池化层、可能涉及通道间计算的逐像素操作，其中包含逐元素操作、堆叠运算符以及 1×1 卷积层的变体。
我们基于四条标准来设计搜索基元。
- 首先，它们需要能够高效地进行评估。
- 其次，我们需要能够轻松地将它们组合成有意义的处理流程。
- 第三，它们需要是可微的，以便能够进行端到端的训练。
- 最后，它们需要足以表达现有的去马赛克算法 ------ 既包括经典的前馈式去马赛克算法，也包括基于深度学习的去马赛克算法和超分辨率算法（图 3）。
由于在马赛克图案中绿色通道具有更高的采样率，受经典去马赛克算法的启发，我们对程序搜索进行分解，首先重建绿色图像，然后利用绿色图像来引导另外两个通道的重建。经典去马赛克算法通常使用一种通过从若干局部方向滤波器中进行边缘自适应选择来重建绿色通道。
与传统的神经架构搜索（NAS）不同，我们同时纳入了一维和二维分组卷积层，以便实现方向滤波器功能，还纳入了 softmax 基元、逐元素乘法以及求和归约操作，从而能够在不同滤波器之间进行可微选择。一维卷积相较于线性可分的二维卷积还能节省计算量。我们还纳入了 1×1 卷积层的变体以及逐通道求和归约操作，以促进通道间的信息传递。
经典的去马赛克算法通常通过预测红、蓝通道与绿色通道的差值来重建红、蓝通道。我们纳入了诸如加法和减法之类的逐元素运算来模拟这一技术。我们的逐元素运算符以及堆叠运算符还通过沿通道维度组合中间输出来支持残差连接。
一些基于深度学习的去马赛克方法会使用一个打包层，该层将拜耳模式中每个重复的 2×2 网格打包成四个平移不变的图像通道（两个绿色、一个蓝色和一个红色），还会使用一个解包层，将一个 4 通道图像转换回全分辨率网格。但马赛克图案可能具有不同的周期性。例如，拜耳模式的周期是 2×2，而 X-Trans 模式的周期是 6×6。我们通过按比例参数化的打包和解包操作来支持对不同马赛克图案的正确处理。解包层也常被深度超分辨率模型用于上采样层，以便将信息从通道维度转移到空间维度。
上采样和下采样是针对去马赛克、超分辨率等任务的另一组重要的特定领域构建模块。与添加卷积层的传统方法相比，下采样是一种有效的替代方案，它可用于扩大模型的感受野，以便识别高频模式并避免摩尔纹伪影。我们支持多种类型的上采样，并允许搜索来选择使用哪些类型。我们的上采样操作符包括：解包、双三次插值以及通过转置卷积实现的学习型上采样。对于下采样，我们使用打包操作或者学习型下采样，学习型下采样是一种步长卷积，其滤波器宽度是步长的两倍，以实现适当的抗锯齿效果。
最后，我们设计了插入函数，这些函数会将预测的颜色值插入到对应的颜色通道中，并结合马赛克图像中给定的颜色值，从而生成完整的红、绿、蓝图像。这些专门的函数使我们的程序能够进行位置感知型的颜色预测，并且通过仅预测缺失值来节省计算量。
总体而言，我们使用了以下构建模块：
- 逐元素运算：加法（Add）、减法（Sub）、乘法（Mul）、堆叠（Stack）、修正线性单元（ReLU）；
- 卷积运算：一维分组卷积（GroupedConv1D）、二维分组卷积（GroupedConv2D）；
- 滤波器选择：SoftMax；
- 逐像素通道间计算：1×1 分组卷积（GroupedConv1x1）、分组求和（GroupedSum）、交错求和（InterleavedSum）；
- 上采样：解包（Unpack）、双三次插值（Bicubic）、学习型上采样（LearnedUpsample）；
- 下采样：打包（Pack）、学习型下采样（LearnedDownsample）；
- 插入：绿色插入（GreenInsert）、色度插入（ChromaInsert）。
卷积运算符一维分组卷积（GroupedConv1D）、二维分组卷积（GroupedConv2D）以及 1×1 分组卷积（GroupedConv1x1）是由分组数量和输出通道数量来进行参数化的。除了 1×1 卷积的情况外，滤波器宽度均为 3。SoftMax 在使用时不带有缩放参数。分组求和（GroupedSum）和交错求和（InterleavedSum）都会接收一个 N N N 通道图像，并通过对通道组内的值进行求和来生成一个 N / k N/k N/k 通道图像。分组求和是对连续的通道进行求和，而交错求和则是通过对通道 i + j N / k i + jN/k i+jN/k（其中 0 ≤ j < k）进行求和来输出通道 i。所有操作均以单精度浮点数形式进行。
Search Algorithm
- 图 4 我们搜索的输入和输出：左边展示的是种子程序，它们源自先前的相关工作，右边展示的是通过我们的搜索找到的四个绿色预测模型，按照计算成本和质量递增的顺序排列。我们使用去马赛克网络（Demosaicnet）的一个变体以及一个受自适应同态滤波（AHD）启发的多分辨率去马赛克程序作为我们的种子模型。请注意，这些生成的多样化模型与原始种子模型有显著差异，随着质量的提升，它们的有向无环图（DAG）结构变得越来越复杂，而且它们是如何将传统的神经构建模块与我们特定领域的原语相结合的。（卷积层在其滤波器尺寸后标注了分组数量）
我们将多目标遗传搜索和梯度下降相结合，以找到呈现为上述构建模块有向无环图（DAG）形式的具有帕累托优势的处理流程。我们的搜索空间需要特定的变异规则，这既是为了处理复杂的有向无环图结构（在这种结构中，节点可能有多个下游父节点以共享计算），也是为了遵循特殊运算符的约束条件（例如，某一阶段的所有输入都要上采样 / 下采样到相同分辨率，并且输出需要达到正确的目标分辨率）。
由于我们的搜索空间非常大，并且允许任意的有向无环图存在，我们利用特定领域知识对其进行分解，以降低其组合复杂性，利用剪枝规则排除明显不佳或低效的程序，并使用快速且稳健的训练方法来快速探索数千个模型。
我们搜索的目标是探索快速去马赛克算法的设计空间，因此我们用两个高效的模型来对其进行初始化，如图 4 所示，这两个模型是基于先前的最先进设计构建的。
其中一个种子模型是去马赛克网络（Demosaicnet）的一个小型简化版本，其层数和通道数量都更少。我们没有将已发布的全尺寸去马赛克网络用作种子模型，是因为它的吞吐量远远低于（几乎低了两个数量级）高效去马赛克算法的吞吐量范围。
我们的另一个种子模型的灵感源自梯度 Halide 模型（GradientHalide model），不过它通过上采样和下采样运算符使用了多种分辨率。
为了在一系列计算成本范围内优化程序质量，我们将程序划分到不同的成本层级中。在每一轮搜索生成之后，我们会在每个成本层级中保留排名前 20 的程序，并从每个层级中抽取 12 个模型进行变异以用于下一代。随着程序规模变大，每个层级区间的宽度会翻倍。成本层级区间起始于绿色插值程序的 0 - 200 次浮点运算以及蓝 / 红插值程序的 0 - 300 次浮点运算。蓝 / 红插值程序允许有更高的成本，因为它们会使用一个绿色插值程序来生成自身的一个输入。
在每一代中，新变异生成的程序会在一个小数据集上进行几个轮次的训练。在搜索过程中，程序会在去马赛克网络（Demosaicnet）数据集的 10 万张图像子集上训练 6 个轮次。人们可以增加成本层级的数量和宽度，以覆盖吞吐量更低的程序。不过，这样做会因每一代中需要训练更多成本更高的程序而增加搜索时间。
成本估算：我们通过遍历有向无环图（DAG）并统计每个节点所执行计算所需的浮点运算次数来估算每个程序的成本，将那些可以使用融合乘加（FMA）执行的运算视为单次运算。一个程序的总成本就是各节点成本之和。尽管这种衡量方式忽略了数据移动，但正如第 4 节所展示的那样，它能够成功估算相对运行时间，因为我们的程序被限制在那些可以为了提高局部性而进行高效分块或融合的运算操作范围内。
在每个成本层级内，我们使用重要性抽样来选择哪些程序要进行变异并保留至下一代。与层级内的其他程序相比，选择某个程序的可能性与其帕累托排名成反比。处于帕累托前沿的模型排名为 1。在移除所有排名为 1 的模型后将会处于前沿的那些模型排名为 2，依此类推。因此，模型是基于它们与当前帕累托前沿的接近程度来进行重要性抽样的。
Mutations 变异操作
每个被选中用于变异的程序会通过以下操作之一来进行修改：
- 插入操作：从构建模块中随机均匀地选择一个构建模块，并将其添加到有向无环图（DAG）中程序的一个随机位置上。如果所选的插入位置对于该构建模块而言是无效的，那么这个变异规则会继续随机选择位置，直至找到一个有效位置为止。如果要插入的操作有两个操作数，那么第二个操作数会被选定为该程序中某个现有子有向无环图（sub-DAG）的输出。这会引入循环连接和共享计算。
- 删除操作：从程序中随机均匀地选择一个有向无环图（DAG）节点并将其移除。
- 分辨率变更：插入、移除以及移动上采样和下采样操作是比较棘手的。我们必须确保所有改变分辨率的变异操作都能保留正确的输出分辨率，并且不允许将具有不同分辨率的中间输出进行合并。为解决这一问题，我们使用图着色的方法，按照每个节点的输出分辨率为其上色。这种表示方式使得我们能够通过操控颜色边界轻松地执行改变分辨率的变异操作。
-- 创建分辨率变更：从一个颜色边界内随机选择一个子图，使其以相对于其父边界更低的分辨率进行计算。分辨率缩放因子会随机选择为 2 或 3。
-- 移除分辨率变更：随机选择一个颜色边界，通过删除或插入上采样和下采样操作来移除它所导致的分辨率变更。
-- 缩小分辨率子图：随机选择边界处的一个节点，并将其分辨率变更为边界另一侧相邻颜色所对应的分辨率。
-- 交换分辨率操作：随机选择一个改变分辨率的节点，并变更其所使用的上采样或下采样函数的类型。 - 解耦操作：随机选择一个被多个使用节点共享的子有向无环图（sub-DAG）并对其进行复制，以便后续的变异操作能够分别修改每个副本。
- 通道数量变更：随机选取一个卷积操作，并修改其输出通道数量。
- 组变更：随机选取一个卷积操作，并修改其通道分组情况，其中允许的分组是正在被变异操作的输入和输出通道数量的任何公因数。我们的搜索发现使用分组卷积能够有效地节省计算量（图 4）。
- 绿色输入变更：仅在搜索红 / 蓝插值程序期间使用，它会修改一个红 / 蓝程序，使其使用不同的绿色子程序作为输入。
搜索更倾向于插入和删除操作而非其他变异操作，因为我们发现这样能使各代程序更快地提高质量并降低成本。
Fast Training
每个经过变异的程序都会使用梯度下降法在一个包含高难度马赛克图像块的数据集上进行训练。理想情况下，我们本应能够在整个数据集上对程序进行训练直至收敛；然而，训练时间（每个程序需要 5-10 个 GPU 小时）会使搜索变得难以进行。因此，我们采用了一种策略，既能大幅缩短训练时间，又能对程序质量进行较好的预估。
首先，我们将训练限定在使用 260 万张图像中固定的 10 万张随机图像子集上，并且只训练 6 个轮次。这两个参数都是基于对程序空间训练动态的探索而设定的。我们发现这样的设置能够对各选项进行良好的排序，其结果与经过完整训练的结果相当接近，但每个程序仅需 10 - 15 个 GPU 分钟。
其次，为避免因灾难性权重初始化而低估程序质量，我们会初始化三组权重。我们同时对这三组权重进行一个轮次的训练。在一个轮次之后，我们继续使用具有最佳验证峰值信噪比（PSNR）的那组权重进行训练。我们观察到，通过比较不同初始化情况下的性能表现，通常可以在第一个轮次内检测出不良的权重初始化情况。这种策略在训练时间仅有小幅增加的情况下，极大地提高了我们质量评估的稳健性。
Efficiently Structured Search：
我们通过多种方式构建搜索空间，以提高在固定时间内于质量-效率权衡空间中找到有价值点的可能性。由于训练每个程序的成本颇高，除了构建模块的结构之外，我们还采用了三个关键策略来引导搜索朝着有用的程序方向进行：
- 在搜索过程中，我们将变异选择倾向于插入和删除操作；
- 我们会在训练程序之前就尽早舍弃那些可能较差的程序；
- 我们将搜索分解到估计绿色值的子程序空间以及利用这些绿色值来插值计算红色和蓝色值的子程序空间中。
由于我们的搜索允许考虑由我们的构建模块组成的所有可能的有向无环图（DAG），变异操作可能会产生一些程序，这些程序不太可能在计算量和去马赛克结果质量之间的权衡中体现出有价值的点。例如，一次变异可能会插入一个与其前面操作互为逆运算的操作（比如先加法后减法，或者先下采样后上采样）。同样，插入操作也可能导致产生无意义的运算（例如，对单个通道进行交错求和）。为避免在这些程序上浪费宝贵的训练时间，我们会在训练阶段之前就尽早舍弃它们。
我们注意到，先前一些关于去马赛克算法的研究工作，由于已知的绿色值数量是其他颜色值数量的两倍，在对绿色值进行插值时采用了与对蓝色和红色值插值不同的策略。而且绿色值会首先被预测，因为它们对于预测红色和蓝色值起着有用的引导作用。因此，我们通过将搜索分解为两个独立的搜索空间来降低搜索的组合复杂性：一个是用于绿色预测程序的搜索空间，另一个是用于红色和蓝色预测程序的搜索空间。我们首先对绿色插值程序进行搜索，并选择一组在一系列计算成本范围内具有帕累托优势的生成程序。然后，在进行红 - 蓝搜索时，可以选择这些具有帕累托优势的绿色插值程序中的任何一个，将其作为马赛克图像的额外输入，以帮助引导红色和蓝色的插值操作。我们在 4.2 节中表明，这种策略比搜索一个联合插值所有颜色值的单一程序的替代方法更具效率。
Compiling Programs to Optimized Implementations
将我们的程序编译为高效的单指令多数据（SIMD）代码是简单且完全自动化的，因为我们在选择构建模块时就特意考虑到了这一点。我们通过将程序转换为 Halide 代码来进行编译。我们遍历模型图，将每个节点映射到其对应的 Halide 实现代码上，通常每个节点对应的代码为 5 - 10 行。
Halide 代码还需要一个 "调度计划"，用于指定算法在硬件上的运行方式。同样由于我们对搜索空间的选择，生成这个调度计划也是简单且完全自动化的。我们所有的节点在空间上的支持范围都较小，所以我们以分块的方式对整个处理流程进行融合。我们根据输出的每个分块按需计算各个阶段。只有一个使用节点的节点，只要不会导致重复计算（即当使用节点不是基于通道的模板运算或归约运算时），就会内联到该使用节点中。执行跨输出通道变化计算的节点（例如，堆叠操作符 Stack）会在输出通道上展开。所有跨输入通道进行归约运算的节点（例如，卷积层）会在输入通道上完全展开。上采样和下采样操作从不内联，以避免使用阶段的寻址模式变得复杂。