背景

众所周知，在传统图像处理领域，不同的处理任务都有各自的一套评价标准，比如降噪任务中的 PSNR 和 SSIM，自动白平衡中的角度差，颜色还原中的色差，等等。但是，当我们把应用场景切换到 CV 领域后，照搬这套传统的图像质量评价（Image Quality Assessment, IQA）体系来评价一个成像系统的性能，很有可能无法得到客观、公正的结果。

举个例子，在目标检测任务中，假设一个检测模型的训练样本全部来自于一台 360P 的古董监控摄像头，那么在部署这个模型时，即使喂给它一张 4K HDR 图像，其检测精度往往可能还不如另一张同样来自这个 360P 摄像头的图像——并不是说 4K HDR 图像的质量不好，但是在当前这一场景下，它与已经训练好的数据驱动的 deep learning 模型并不适配。因此，在 CV 任务的场景下，当一张图像的观察主体从人眼+人脑的组合切换为某个 DL 模型时，我们有必要对原有的图像质量评价体系进行重新定义。

在针对 CV 任务提出的 IQA 方案（以下称为 CV-IQA）中，用户注重的不再是解析度、信噪比、颜色准确性这些面向人眼的图像质量指标，而是会更加关注于物理世界中的信号在数字域中的可复现性（Reproducibility）和复现可靠性（Reproduction Reliability）。

以下图为例，左边是一块限速牌在物理世界中真实的亮度分布（用单位面积辐射出的光子数表示），右边是使用某台相机拍摄下的该限速牌的图像。由于任何光学系统和半导体器件都存在一些固有缺陷，因此在最终的图像中，限速牌各个部分的对比度，或者说信号的可观测差异，要比物理世界中的真实限速牌差得多。

差得多，那么究竟差了多少呢？我们试图从传统 IQA 指标中去寻找一个能够表征这种差异程度的指标：PSNR？量纲都对不上，PSNR 肯定没法计算，而且实际环境下这两张图也很难实现像素级的对齐；对比度？右边这张图像里有那么多的像素，要计算谁和谁之间的对比度？想了一圈发现，在这种实验设置下，传统指标似乎总是差了那么一点意思。

什么是 CDP

CDP（Contrast Detection Probability）指标是 IEEE P2020 小组（Automotive Image Quality Working Group）针对自动驾驶场景提出的一套 CV-IQA 方案，旨在克服传统 IQA 方案（如 IEEE Std 1858、EMVA1288、ISO 12233）用于车载影像质量评价时存在的一些缺陷。

用一句话概括 CDP 存在的意义：CDP 表征了成像系统对物理世界辐射信号差异的复现能力（原文：CDP is a metric to describe the performance of an imaging system to reproduce contrasts in the physical scenes）。

划两个重点：

第一，在 CV 任务中，我们希望成像系统能够重点关注物理信号间的差异，而非信号本身的绝对大小。例如对于物理世界里亮度为 $(500\mathrm{cd/m^2}, 600\mathrm{cd/m^2})$ 的两个物体，我们关心的是在最终的输出图像中这两个物体能否进行区分，而并不关心具体的灰度值大小，它们可以是 $(200, 255)$，也可以是 $(50, 64)$。这一点很好理解，因为一旦这两个物体在图像中无法区分，其梯度便退化为零，基于特征提取的 CV 模型自然无法从中提取出任何差异化信息，不管是分类还是检测还是分割任务，统统歇菜。

第二，CDP 考量的是成像系统能否对信号差异进行准确复现。所谓准确复现，即是说物理世界里一组信号之间的差异有多大，在最终输出图像中，这个差异应该还是有多大。信号的差异在经过成像系统前后需要维持在一个相同的水平，既不能放大，也不能减小。还是上面那个例子，假如我们用 Weber 对比度作为计算信号差异的指标，那么物理世界中这两个信号之间的差异为 $K_{\mathrm{world}}=(\frac{600}{500}-1)=0.2$。如果有两个成像系统 A 和 B，在它们的输出图像 $\mathcal{I}_A$ 和 $\mathcal{I}_B$ 中，这两个物体的灰度值分别为 $(200, 255)$ 和 $(55, 64)$，那么有：

这时我们就说成像系统 A 有着更好的信号差异复现能力，因为 0.275 相比 0.28 更接近物理世界中的真实对比度值——0.2。

从直觉上来说，我们往往会希望成像系统能够把真实世界中的信号差异进行一定程度的放大，比如 ISP 里的边缘增强、对比度增强等模块，都是为了突出和强调信号间的差异。但是在 CV 的语境下，我们希望成像系统能够尽量做到信息（即信号差异）的忠实还原，因此 B 系统虽然对信号差异进行了放大，但是其 CDP 指标却要低于 A 系统。这个例子也直观体现了 CV-IQA 与传统 IQA 之间的设计思想差异。

铺垫了这么长，那么为什么 CDP 这个指标可以针对自动驾驶场景下（或者说得更泛一点，大部分 CV 场景下）成像系统的质量进行评价呢？我个人的理解是，图像不同区域之间信号的可分辨性（distinguishability）是 CV 模型执行正确计算的必要条件，因此可以用 CDP 来衡量一个成像系统对物理世界中信号的检测能力的上限——如果一个成像系统的 CDP 很高，其检测能力未必很好，但是如果 CDP 很低，其检测能力一定不好。换句话说，CDP 并不是一个万能的 CV-IQA 方案，它只是为我们提供了一个相对客观的评价维度，同时补足了传统 IQA 中存在的一些短板。

技术细节

这一节重点介绍 CDP 这个指标究竟如何对成像系统的质量进行定量评价，以及给定一台相机，我们要如何计算出最终的 CDP 结果。

继续沿用上节中的例子。

假设物理世界中有两块亮度分别为 $500\mathrm{cd/m^2}$ 和 $600\mathrm{cd/m^2}$ 的理想漫反射平面：

上一节中我们已经计算出了它们之间的 Weber 对比度为 $K_\mathrm{world}=0.2$。同时，使用它们的亮度平均值 $\frac{500+600}{2}=550\mathrm{cd/m^2}$ 作为这一组信号的物理亮度 $L$。

现在我们使用一台相机分别去拍摄这两块平面，假设得到这样两幅图像：

从两张图像中各自随机抽取一个像素，如下红色小方框所示。不妨假设它们在图像中的灰度值分别为160和206。

CDP 要求我们在计算这两个像素的对比度之前，先把它们从图像灰度值域（digital-number domain）映射回物理亮度域（luminance domain）。显然，这时候我们并不清楚160和206的灰度值到底对应多少 $\mathrm{cd/m^2}$ 的物理亮度值，因此在执行这步操作之前，我们需要为当前成像系统构建一个输出 $\rightarrow$ 输入的反映射函数（inverse system function）。

P2020 中并没有明确说明如何构建这个反映射函数，我个人采用的方法是，先在横坐标为物理亮度、纵坐标为灰度值的平面内绘制出输入 $\rightarrow$ 输出曲线（即我们平时说的系统响应曲线），然后把当前要计算的灰度值作为 $y$ 坐标，用 bilinear 插值的方式计算出它在曲线中对应的 $x$ 坐标，以此作为该像素对应的物理亮度值。

输入 $\rightarrow$ 输出曲线的绘制可以预先通过标定实验来完成。我们假设当前相机的灰度值与物理亮度之间满足如下近似于 Gamma 曲线的关系：

对于160和206的灰度值，我们在这条曲线上可以内插出其对应的物理亮度大约为 $505\mathrm{cd/m^2}$ 和 $640\mathrm{cd/m^2}$。换句话说，对于这台相机，digital-number domain 中的160 (206) 数值，大约对应到 luminance domain 中的 $505\ (640)\mathrm{cd/m^2}$：

有了 luminance domain 中的一组亮度值之后，我们可以计算它们的 Weber 对比度：

除了可以按照如上两个红色小方框进行像素选取之外，我们一共有 $M\times{}N$ 种方式从两个 RoI 中各自选取一个像素构成一组 pixel-pair，其中 $M$ 和 $N$ 分别代表两个 RoI 中的像素数。如果我们对这 $M\times{}N$ 组 pixel-pair 都计算一次 Weber 对比度，则可以得到一个关于对比度的分布直方图。

两个平面在图像中的 RoI 都是 $32\times{}32$ 的正方形（$M=N=32^2=1,024$），因此我们一共可以得到 $1,024\times{}1,024\,\approx{}10^{6}$ 个对比度值，它们构成的直方图如下所示：

对这个直方图进行归一化，我们就得到了一组 RoI-pair 之间的对比度概率分布函数 $p(K)$：

（到此终于解释了为什么 CDP 中的那个 P 是 probability……）

在这个例子中，物理世界里两个平面之间的信号差异为 $K_\mathrm{world}=0.2$，而由于噪声以及非线性的存在，在成像系统记录下的信号中，它们之间的差异不再是一个具体的数值，而是一组关于对比度的概率分布函数 $p(\cdot)$。

在正式给出 CDP 的计算公式之前，需要先引入置信度的概念。

回想一下上文中我们对 CDP 的定义：CDP 用来衡量成像系统对物理信号差异的复现准确性。对于 $K_\mathrm{world}=0.2$ 的输入信号，理想情况下，我们当然希望成像系统能够分毫不差地复现出这个对比度数值，这时上图中的概率分布理应是一个落在 $K=0.2$ 处的 Dirac 函数。然而在实际系统中，$p(\cdot)$ 必然有一定的带宽，我们也会适当地放松对「准确复现」这一概念的定义，即，对于成像系统复现出的一组信号，只要它们的对比度落在 $\left[K_\mathrm{world}(1-\epsilon), \ K_\mathrm{world}(1+\epsilon)\right]$ 区间之内，我们就将其视为一次「准确复现」。这个 $\epsilon$，正是一个预设的置信度值。

回到上面的例子中，假设我们设定 $\epsilon=10\%$，这时对于任意一组 pixel-pair，只要它们的对比度位于 $[0.18,\ 0.22]$ 区间之内，我们都认为这组 pixel-pair 成功地对物理对比度进行了准确复现。在上面的概率分布函数中画出0.18和0.22所在的位置，其包围的区域与 $p(\cdot)$ 取交集得到的面积，相对于 $p(\cdot)$ 与 $K$ 轴包围的面积之比，即为最终的 CDP 值：

根据上图阴影部分占整个曲线下面积（AUC）的比例，CDP 值大约在 30% 左右。

我们注意到公式中的 $\mathrm{CDP}$ 是一个关于 $L$ 的函数，这是因为，这一计算结果仅对平均亮度 $L=\mathrm{550cd/m^2}$ 的一组输入信号成立，当任一平面的物理亮度发生改变时（这时两个平面的平均亮度也会改变），我们又需要重新构建一次 $p(\cdot)$，再计算一个新的 CDP 值。

对于上面的例子，我们使用如下定义对系统的信号复现能力进行描述：

给定 10% 的置信度（$\epsilon=0.1$），在 $550\mathrm{cd/m^2}$ 的亮度水平下，当前成像系统有 30% 的概率准确复现出物理世界中对比度为 $K_\mathrm{world}=0.2$ 的一组信号。

相比于传统的 IQA 指标，CDP 在对成像系统的性能进行评价时把外界的信息也作为了自变量之一，从而实现了被测系统输入（辐射信号）与输出（数字信号）之间的互动。

计算示例

让我们来看一个具体的示例。

假设有216个不同物理亮度的色块，它们当中最亮亮度约为 $5\times{}10^4\mathrm{cd/m^2}$，最暗的约为 $7.5\times{}10^{-2}\mathrm{cd/m^2}$，按下图所示排布。（这一排布方式参考了 Image Engineering 的 DTS 仪器）

现在我们使用一个基于长中短曝光三帧合成的 HDR 相机对这个场景进行拍摄，得到如下 RAW 图像（实际是 16bit RAW，这里为了能够正常在屏幕上显示所以做了一个16bit $\rightarrow$ 8bit 的线性压缩）：

使用 DOL-HDR 技术的相机在特定的亮度拼接区域会出现明显的 SNR drop，在这个例子中如果我们把图像右上角进行放大，可以看到某些色块对应的噪声强度明显大于其他色块：

我们把图像中每个色块单独提取出来并计算其平均灰度值（这张图里每个色块大约占 $40\times{}40$ 个像素），可以得到216个不同的灰度值。把色块的物理亮度值作为横坐标，把图像灰度值作为纵坐标，可以绘制出这个相机的响应曲线（这是一个仿真的相机，实际相机的曲线当然不可能这么规整 😳 ）：

有了响应曲线，就可以使用上一节提到的插值的方法来构建图像灰度值域到物理亮度域的反映射函数。

我们每次从这216个色块中任意选取两个色块作为一组 RoI-pair，然后按照上一节的方法可以计算得到一个 CDP 值。对于216个色块，一共可以不重复地选出 $C_{216}^{\,2}=23,220$ 组 RoI-pairs，因此最终一共将得到23,220个 CDP 值，以及23,220个物理对比度 $K_\mathrm{world}$。

在实际的应用中，我们通常只关心某些物理对比度 $K_\mathrm{world}$ 下的 CDP 情况，因此通常只会对这23,220组 RoI-pairs 中满足 $K^\ast (1-\delta)\le{}K_\mathrm{world}\le{}K^\ast(1+\delta)$ 的那些个 pairs 进行 CDP 的计算，其中 $K^\ast$ 表示用户指定的目标对比度，$\delta$ 是一个允许的 tolerance，一般取0.1即可。

在这一节的例子中，当我们指定目标对比度 $K^\ast=0.2$、$\delta=0.1$ 时，一共可以得到291对满足 $0.18\le{}K_\mathrm{world}\le{}0.22$ 的 RoI-pairs。如果我们把这291个 CDP 值绘制出来，并以物理亮度 $L$ 作为横坐标，可以得到这么一个分布情况：

这张图完整表征了被测相机对于指定的目标对比度（0.2），在不同环境亮度下对物理信号进行准确复现的能力。

类似地，我们还可以绘制出 $K^\ast=0.06$、$K^\ast=0.1$ 以及 $K^\ast=0.3$ 下的 CDP 分布情况（这三个数值是 Image Engineering 推荐的目标对比度值）：

考虑到一个成像系统往往需要在不同的场景下工作，因此在使用 CDP 对成像系统的性能进行评价时，一般都会给出某个（或某几个）目标对比度下 CDP 关于物理亮度的散点分布，而不是固定一档亮度下的单个 CDP 数值。这一背景也导致了 CDP 无法以一个简单指标的形式向用户进行交付，同时也不容易对两个系统的性能差异进行定量比较（P2020 小组并没有提到如何评判两组 $L-\mathrm{CDP}$ 散点图之间孰优孰劣），这也许也是导致 CDP 始终没有得到普遍应用的原因之一吧。

Reference

• Geese M, Seger U, and Paolillo A, Detection Probabilities: Performance Prediction for Sensors of Autonomous Vehicles. 2018
• Artmann U, Geese M, Gäde M, Contrast Detection Probability – Implementation and Use Cases. 2019
• Imatest Documentation: Contrast Detection Probability. https://www.imatest.com/docs/cdp/
• ISO 12232:2019 Photography — Digital still cameras — Determination of exposure index, ISO speed ratings, standard output sensitivity, and recommended exposure index.
• Robin B. Jenkin, Comparison of Detectability Index and Contrast Detection Probability. 2019
• Lukas Ebbert, Implementierung von CDP: Entwicklung eines Programmiercodes in Python zur Untersuchung und Messung von CDP bei Fahrerassistenzkameras. 2018

背景

最近的一个项目中需要用 PyTorch 实现一套可微分的 Non-local Means (NLM) 降噪算法，并通过样本学习建立一个图像内容与降噪强度参数 $h$ 的映射模型。

在 PyTorch 框架下，实现传统图像处理算法往往有着不止一种的写法，例如，为了实现 $3\times3$ 的均值滤波，我们既可以使用一个归一化的 box filter 对图像进行卷积，也可以对图像进行9次位移（shift/roll）并逐像素求和，最后再将每个像素值除以9。在模型进行正向推断（forward）时，不同的写法在显存开销上或许不存在太大的区别，但是当我们使用 Autograd 进行逐层的梯度反向传播（backward）时，不同的写法往往对应了截然不同的显存占用。以 NLM 为例，在我的实验中，对于同一张 $512\times{}512$ 的输入图像，不同的 tensor 操作写法可以带来180%以上的训练显存差异。当然，在代码设计没有明显缺陷的情况下，节省显存必然意味着运算效率的下降，不过对于我的这个项目来说，训练阶段的耗时并不是瓶颈所在，batch size 才是影响模型性能的主要因素，因此利用时间换取空间仍然是一个非常划算的选择。

先简单地回顾一下 Non-local Means 算法：给定一幅噪声图像，对于图像中的每个像素 $x$，以其为中心，在一个 $w\times{}w$ 的搜索窗口（下文代码中以 search_window 命名）内，遍历所有的可能的图像块（一般是一个 $k\times{}k$ 的方形区域，且有 $k<w$，下文代码中以 patch 命名），并计算该图像块与以 $x$ 为中心的那个图像块之间的相似度。

在 $w\times{}w$ 的搜索窗口内，一共存在 $w^2$ 个可重叠的图像块，对于第 $i$ 个图像块，假设其与 $x$ 所在图像块之间的相似度为 $s_i$（注意：是图像块之间的相似度，而不是像素与像素之间的相似度！），且该图像块中心像素的像素值为 $y_i$，那么对于 $x$ 像素，其经过 NLM 降噪后的像素值可以表示为：

其中 $f(\cdot)$ 为一单调递增函数，即，若第 $i$ 个图像块与 $x$ 所在图像块的相似度越高，其在加权平均时所占的比重就越大。实际计算中一般会采用两个图像块之间的欧式距离 $d$ 作为相似度度量，即 $f(\cdot)=\exp\left(-\frac{d}{h}\right)$，其中 $h$ 是一个用于控制降噪强度的超参数。

一点准备工作

为简化起见，假设现在给定了一张噪声图像和一张与之对应的无噪声参考图像（ground-truth），我们希望对该噪声图像求得一个最佳的 $h$ 参数，使之经过 NLM 降噪后与参考图像之间的 PSNR 最大化。

在这个全监督任务中，我们可以直接使用 MSE loss 作为目标函数进行迭代优化：

# train.py

import torch
import torch.nn as nn

from .nlm import NonLocalMeans
from .utils import gen_image_pair


dev = torch.device('cuda')
clean_rgb, noisy_rgb = gen_image_pair('Lenna.png', device=dev, sigma=0.05)

denoiser = NonLocalMeans(h0=0.05).to(dev)
denoiser.train()
optimizer = torch.optim.Adam(params=denoiser.parameters(), lr=0.0001, weight_decay=0.0001)
loss = nn.MSELoss()

for iter_num in range(100):
    denoised_rgb = denoiser(noisy_rgb)
    mse = loss(clean_rgb, denoised_rgb)

    mse.backward()
    optimizer.step()

    psnr = 10 * torch.log10(1.0 / mse)
    print('step{}: PSNR={:.3f} (h={:.5f})'.format(iter_num, psnr.item(), float(denoiser.h.item())))

其中 gen_image_pair 用于生成一对干净图像和噪声图像，为了不影响阅读的连续性，我把它的具体实现放在了文末的附录中；NonLocalMeans 模块接收一个 $N\times{}3\times{}H\times{}W$ 的张量作为输入，返回一个同尺寸的降噪后的张量。NonLocalMeans 的具体实现将放在下文中展开，这里我们只需确保该模块中包含一个可训练的 self.h 参数：

# nlm.py

import torch
import torch.nn as nn


class NonLocalMeans(nn.Module):
    def __init__(self, h0, search_window_size, patch_size):
        super().__init__()
        self.h = nn.Parameter(torch.tensor([float(h0)]), requires_grad=True)
        pass

    def forward(self, rgb):
        pass

NonLocalMeans v1

在第一版的 PyTorch NLM 中，我很自然地想到了空间换时间的方法，即，对于每个像素，将其搜索窗口中的 $w\times{}w$ 个像素拍扁并放置到一个新的维度中。以单通道图像为例，假设输入图像尺寸为 $1\times{}H\times{}W$，搜索窗口宽度 $w=11$，那么可以创建一个尺寸为 $1\times{}H\times{}W\times{}121$ 的张量，其中第 $(1,i,j,k)$ 个元素表示的是在以原图 $(i,j)$ 像素为中心的搜索窗口中第 $k$ 个像素的像素值。

有了这个张量之后，我们就可以将其与原图相减得到一个差值张量（相减之前需要将原图 unsqueeze 出一个新的维度用于 broadcast），并在这个差值张量上利用局部求和+开方的方法得到图像块之间的欧氏距离（出于偷懒的原因，图像块之间相似度的计算仅在亮度通道进行，下同。rgb_to_luminance 的实现见附录）：

# nlm_v1.py

from .utils import rgb_to_luminance, ShiftStack, BoxFilter

EPSILON = 1E-12


class NonLocalMeans(nn.Module):
    def __init__(self, h0, search_window_size=11, patch_size=5):
        super().__init__()
        self.h = nn.Parameter(torch.tensor([float(h0)]), requires_grad=True)

        self.gen_window_stack = ShiftStack(window_size=search_window_size)
        self.box_sum = BoxFilter(window_size=patch_size, reduction='sum')

    def forward(self, rgb):
        y = rgb_to_luminance(rgb)  # (N, 1, H, W)

        rgb_window_stack = self.gen_window_stack(rgb)  # (N, 3, H, W, w*y)
        y_window_stack = self.gen_window_stack(y)  # (N, 1, H, W, w*y)

        distances = torch.sqrt(self.box_sum((y.unsqueeze(-1) - y_window_stack) ** 2))  # (N, 1, H, W, w*y)
        weights = torch.exp(-distances / (torch.relu(self.h) + EPSILON))  # (N, 1, H, W, w*y)

        denoised_rgb = (weights * rgb_window_stack).sum(dim=-1) / weights.sum(dim=-1)  # (N, 3, H, W)

        return torch.clamp(denoised_rgb, 0, 1)  # (N, 3, H, W)

其中的 self.gen_window_stack 方法用来将搜索窗口内的像素拍扁并放到一个新的维度中（ShiftStack 与 nn.Unfold 类似，但是其实现比 Unfold 更省显存）；self.box_sum 方法用来计算 $k\times{}k$ 窗口内的各像素值之和。它们的具体实现同样放在了文末附录中。

让我们先看下迭代过程中的 PSNR 和 $h$ 参数的变化，确保训练过程可以正常收敛：

step0: PSNR=28.600 (h=0.05050)
step1: PSNR=28.675 (h=0.05098)
step2: PSNR=28.747 (h=0.05146)
step3: PSNR=28.818 (h=0.05194)
...
step97: PSNR=32.882 (h=0.10972)
step98: PSNR=32.888 (h=0.11034)
step99: PSNR=32.894 (h=0.11096)

再看下降噪效果，确保这段代码能够正常实现 NLM 算法：

在这个版本的 NLM 中，对于单张 $512\times{}512$ 的 RGB 图像，一次 forward 大概需要占用 2.3G 显存（已使用 with torch.no_grad() 关闭梯度计算），且训练阶段（forward+backward，见上文第一段代码块）的显存占用与 forward 基本持平：

forward:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 106...  Off  | 00000000:0D:00.0  On |                  N/A |
| 39%   57C    P2    90W / 120W |   2305MiB /  6069MiB |    100%      Default |
+-------------------------------+----------------------+----------------------+

forward + backward:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 106...  Off  | 00000000:0D:00.0  On |                  N/A |
| 43%   59C    P2    96W / 120W |   2307MiB /  6069MiB |    100%      Default |
+-------------------------------+----------------------+----------------------+

显然，这个版本中，中间变量 rgb_window_stack、y_window_stack、distances、weights 均是5维张量，当使用默认的搜索窗口宽度 $w=11$ 时，为了存下它们，PyTorch 至少需要额外申请121倍图像大小的显存。

NonLocalMeans v2

在 NLM 算法中，用来计算相似度的图像块的尺寸 $k$ 通常要小于搜索窗口的尺寸 $w$，因此一个很容易想到的显存优化策略就是，将额外多出来的那个维度用来存储图像块内的各个像素值，而非搜索窗口内的各个像素值：

# nlm_v2.py

class NonLocalMeans(nn.Module):
    def __init__(self, h0, search_window_size=11, patch_size=5):
        super().__init__()
        self.h = nn.Parameter(torch.tensor([float(h0)]), requires_grad=True)

        self.gen_patch_stack = ShiftStack(window_size=patch_size)
        self.r = search_window_size // 2

    def forward(self, rgb):
        batch_size, _, height, width = rgb.shape
        weights = torch.zeros((batch_size, 1, height, width)).float().to(rgb.device)  # (N, 1, H, W)
        denoised_rgb = torch.zeros_like(rgb)  # (N, 3, H, W)

        y = rgb_to_luminance(rgb)  # (N, 1, H, W)

        y_patch_stack = self.gen_patch_stack(y)  # (N, 1, H, W, k*k)

        for x_shift in range(-self.r, self.r + 1):
            for y_shift in range(-self.r, self.r + 1):
                shifted_rgb = torch.roll(rgb, shifts=(y_shift, x_shift), dims=(2, 3))  # (N, 3, H, W)

                shifted_y_patch_stack = \
                    torch.roll(y_patch_stack, (y_shift, x_shift), dims=(2, 3))  # (N, 1, H, W, k*k)

                distance = torch.sqrt(((y_patch_stack - shifted_y_patch_stack) ** 2).sum(dim=-1))  # (N, 1, H, W)
                weight = torch.exp(-distance / (torch.relu(self.h) + EPSILON))  # (N, 1, H, W)

                denoised_rgb += shifted_rgb * weight  # (N, 3, H, W)
                weights += weight  # (N, 1, H, W)

        return torch.clamp(denoised_rgb / weights, 0, 1)  # (N, 3, H, W)

这个版本中，我使用两层循环完成对搜索窗口中每个像素的遍历，因此无论搜索窗口设定得多大，其影响的只是运行速度，而非显存。当图像块宽度 $k=5$，搜索窗口宽度 $w=11$ 时，forward 阶段显存开销理论上可以降低至 v1 版本的 $\frac{25}{121}\approx{}21\%$。

实测发现，对于同样尺寸的输入图像，这个版本的 NLM 一次 forward 只需要占用大约 600M 显存，而训练阶段则是 1.2G：

forward:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 106...  Off  | 00000000:0D:00.0  On |                  N/A |
| 45%   61C    P2    92W / 120W |    603MiB /  6069MiB |    100%      Default |
+-------------------------------+----------------------+----------------------+

forward + backward:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 106...  Off  | 00000000:0D:00.0  On |                  N/A |
| 40%   57C    P2    96W / 120W |   1185MiB /  6069MiB |     99%      Default |
+-------------------------------+----------------------+----------------------+

NonLocalMeans v3

再极端一点，如果我们连用来存储图像块内的各个像素值的中间变量都不想创建呢？显然，可以在 v2 版本的基础上，再加两层循环完成图像块相似度计算（BoxFilter 实际上也是利用两层循环来实现卷积的效果），这个时候，没有任何新的维度需要添加到中间过程的张量中，所有操作都可以原址完成：

# nlm_v3.py

class NonLocalMeans(nn.Module):
    def __init__(self, h0, search_window_size=11, patch_size=5):
        super().__init__()
        self.h = nn.Parameter(torch.tensor([float(h0)]), requires_grad=True)

        self.box_sum = BoxFilter(window_size=patch_size, reduction='sum')
        self.r = search_window_size // 2

    def forward(self, rgb):
        batch_size, _, height, width = rgb.shape
        weights = torch.zeros((batch_size, 1, height, width)).float().to(rgb.device)  # (N, 1, H, W)
        denoised_rgb = torch.zeros_like(rgb)  # (N, 3, H, W)

        y = rgb_to_luminance(rgb)  # (N, 1, H, W)

        for x_shift in range(-self.r, self.r + 1):
            for y_shift in range(-self.r, self.r + 1):
                shifted_rgb = torch.roll(rgb, shifts=(y_shift, x_shift), dims=(2, 3))  # (N, 3, H, W)
                shifted_y = torch.roll(y, shifts=(y_shift, x_shift), dims=(2, 3))  # (N, 1, H, W)

                distance = torch.sqrt(self.box_sum((y - shifted_y) ** 2))  # (N, 1, H, W)
                weight = torch.exp(-distance / (torch.relu(self.h) + EPSILON))  # (N, 1, H, W)

                denoised_rgb += shifted_rgb * weight  # (N, 3, H, W)
                weights += weight  # (N, 1, H, W)

        return torch.clamp(denoised_rgb / weights, 0, 1)  # (N, 3, H, W)

在这个版本里，一次 forward 和一次 backward 分别只占用了大约 470M 和 1.1G 显存：

forward:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 106...  Off  | 00000000:0D:00.0  On |                  N/A |
| 43%   60C    P2    92W / 120W |    473MiB /  6069MiB |    100%      Default |
+-------------------------------+----------------------+----------------------+

forward + backward:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 106...  Off  | 00000000:0D:00.0  On |                  N/A |
| 43%   59C    P2    72W / 120W |   1113MiB /  6069MiB |     99%      Default |
+-------------------------------+----------------------+----------------------+

这也是我能想到的针对 Non-local Means 算法最显存友好的一种实现了。

Reference

• Wikipedia: Non-local means
• Gist: etienne87/torch_bilateral_gray.py

附录

gen_image_pair

# utils.py

import numpy as np
import skimage.io
import torch
import torch.nn as nn


def gen_image_pair(image_path, device, sigma):
    # :return: two tensors with shape (1, 3, H, W) in [0, 1] range
    clean_rgb = skimage.io.imread(image_path).astype(np.float32) / 255.0
    # additive white gaussian noise
    awgn = np.random.normal(0, scale=sigma, size=rgb.shape).astype(np.float32)
    noisy_rgb = np.clip(clean_rgb + awgn, 0, 1)

    clean_rgb = torch.from_numpy(rgb.transpose(2, 0, 1)).unsqueeze(0).to(device)
    noisy_rgb = torch.from_numpy(noisy_rgb.transpose(2, 0, 1)).unsqueeze(0).to(device)
    return clean_rgb, noisy_rgb

rgb_to_luminance

# utils.py
    
def rgb_to_luminance(rgb_tensor):
    # :param rgb_tensor: torch.Tensor(N, 3, H, W, ...) in [0, 1] range
    # :return: torch.Tensor(N, 1, H, W, ...) in [0, 1] range
    assert rgb_tensor.min() >= 0.0 and rgb_tensor.max() <= 1.0
    return 0.299 * rgb_tensor[:, :1, ...] + 0.587 * rgb_tensor[:, 1:2, ...] + 0.114 * rgb_tensor[:, 2:, ...]

ShiftStack

# utils.py
    
class ShiftStack(nn.Module):
    """
    Shift n-dim tensor in a local window and generate a stacked
    (n+1)-dim tensor with shape (*orig_shapes, w*y), where wx
    and wy are width and height of the window
    """
    def __init__(self, window_size):
        # :param window_size: Int or Tuple(Int, Int) in (win_width, win_height) order
        super().__init__()
        wx, wy = window_size if isinstance(window_size, (list, tuple)) else (window_size, window_size)
        assert wx % 2 == 1 and wy % 2 == 1, 'window size must be odd'
        self.rx, self.ry = wx // 2, wy // 2

    def forward(self, tensor):
        # :param tensor: torch.Tensor(N, C, H, W, ...)
        # :return: torch.Tensor(N, C, H, W, ..., w*y)
        shifted_tensors = []
        for x_shift in range(-self.rx, self.rx + 1):
            for y_shift in range(-self.ry, self.ry + 1):
                shifted_tensors.append(
                    torch.roll(tensor, shifts=(y_shift, x_shift), dims=(2, 3))
                )

        return torch.stack(shifted_tensors, dim=-1)

BoxFilter

# utils.py
    
class BoxFilter(nn.Module):
    def __init__(self, window_size, reduction='mean'):
        # :param window_size: Int or Tuple(Int, Int) in (win_width, win_height) order
        # :param reduction: 'mean' | 'sum'
        super().__init__()
        wx, wy = window_size if isinstance(window_size, (list, tuple)) else (window_size, window_size)
        assert wx % 2 == 1 and wy % 2 == 1, 'window size must be odd'
        self.rx, self.ry = wx // 2, wy // 2
        self.area = wx * wy
        self.reduction = reduction

    def forward(self, tensor):
        # :param tensor: torch.Tensor(N, C, H, W, ...)
        # :return: torch.Tensor(N, C, H, W, ...)
        local_sum = torch.zeros_like(tensor)
        for x_shift in range(-self.rx, self.rx + 1):
            for y_shift in range(-self.ry, self.ry + 1):
                local_sum += torch.roll(tensor, shifts=(y_shift, x_shift), dims=(2, 3))

        return local_sum if self.reduction == 'sum' else local_sum / self.area

从13年刚开始搭建博客的时候我就在关注中文的 webfont 解决方案了，可惜那时候几乎找不到一套成熟的跨平台字体设置方案，Google Fonts 和 Typekit 也远没有现在这么完善，Windows 下微软雅黑丑出天际，我们能做的无非就是老老实实设置好 fallback 来确保不同操作系统下都能显示出我们希望显示的字体。

到了2019年就不一样啦，现在的 Google Fonts 和 Adobe Typekit 都提供了傻瓜式的第三方字体托管，同时 CDN 的普及也确保了不同地域的用户都能够（比较）迅速地加载非本地字体。更关键的是，今天即使是移动设备的带宽也允许用户在1–2秒内完成一套中文字体的加载。我们已经完全有理由抛弃微软雅黑而采用更灵活的中文网页字体设置方案。

font-family

font-family 的填写原则在网络上可以搜出无数文章，这里我仅仅简单说一下我自己的一些总结：

• 英文字体在前，中文字体在后。原因很简单，中文字体包中通常包含了英文字符，反之则不成立，如果先设置了中文字体，那么英文字体根本轮不到 fallback
• 操作系统方面，Apple 系列（macOS + iOS）字体在前（直接用 -apple-system 即可，见下文），Windows 字体在后。这里解释了原因，即，macOS 下的用户有一定概率安装了 Windows 字体（因为安装了 Office），而 Windows 下安装 macOS 字体的用户却少之又少，所以如果 Windows 在前，那么在 macOS 下很有可能就出现「放着苹方不用却用微软雅黑」这种坑爹情况
• 根据访问网站的用户群体不同，可以考虑放弃对于老旧操作系统和浏览器的支持，比如我在使用思源宋体的时候，仅提供了 .woff 和 .woff2 格式的字体文件，完全放弃了老版本的 IE 和 Safari
• 仅对 Windows 用户使用自定义字体，比如思源系列，macOS、iOS、Linux 就使用系统自带字体就好。原因嘛当然是因为 Windows 自带的中文字体实在找不到好看的，而主动安装字体的用户毕竟还是少数（嗯这条是私货 🙄 ）
• 衬线还是非衬线完全看个人审美喜好，比如我就不觉得思源黑体系列好看（所以你现在看到的思源宋体，如果你是 Windows 用户的话）。同时，中文衬线搭配英文非衬线，或者反过来，中文非衬线搭配英文衬线，也完全没有问题

我目前使用两套 font-family 设置。正文（Main text）部分是：

font-family: Georgia, -apple-system, 'Nimbus Roman No9 L', 'PingFang SC', 'Hiragino Sans GB', 'Noto Serif SC', 'Microsoft Yahei', 'WenQuanYi Micro Hei', 'ST Heiti', sans-serif

标题（Headings）和页面导航栏（Navigation）部分是：

font-family: 'Josefin Sans', -apple-system, 'PingFang SC', 'Hiragino Sans GB', 'Microsoft Yahei', 'WenQuanYi Micro Hei', 'ST Heiti', sans-serif

其中 Georgia、-apple-system、'Nimbus Roman No9 L' 分别对应 Windows、macOS/iOS、Linux 下三种系统内置英文字体（其实 Georgia 就已经覆盖三种系统了，后面两个只是出于保险起见），'PingFang SC'（苹方）、'Hiragino Sans GB'（冬青黑体）、'Microsoft Yahei'（微软雅黑）对应三种系统内置中文字体（按照上文说法，Apple 系列在前，Windows 在最后），'Josefin Sans' 和 'Noto Serif SC'（思源宋体）则是两套允许免费使用的第三方字体，已经被我存储在本地。

理论上来说，一旦在 font-family 中放置了类似思源宋体这种第三方字体，那么浏览器就一定会加载该字体文件，因此也就没有必要再考虑后续的 fallback 字体。但是就像在 font-display 里面说到的，为了避免浏览器在完成加载第三方字体文件之前显示一片空白，我们仍然有必要设置 fallback 字体作为备胎。

Google Fonts 的本地化

在 Google Fonts 的官方仓库中，中文字体包通常会被拆分为几十上百不等的若干个字体文件，从而确保打开某个网页时仅加载该页面中需要的字符。这种思路和 font-spider 这类的字体压缩工具类似，但是把加载哪些字体文件完全交给 CSS 判断，因此压缩比自然远不能与后者相比。除此之外，还有下下几点因素让我决定放弃使用 Google Fonts 进行字体在线托管，而选择将其部署在本地：

• 以 font-weight: 400 的思源宋体为例（Google Fonts 里的名字叫做 Noto Serif SC），对于单个阅读时间在3分钟左右的中文页面，需要加载18–22个左右的 .woff2 字体文件，总的体积在 900KB–1.1MB 左右。而如果选择将完整的字体包存储在本地的话，也仅需要 2MB 不到的空间而已。在我看来，用请求数的大幅增加（20 vs 1）来换取体积的减少（1MB vs 2MB）似乎并不怎么划算
• 将字体数据分散在20来个字体文件中，并为每个文件都设置 font-display: swap 属性（Google Fonts 的默认策略，后文会解释），会导致整个页面由于字体的切换而产生一种不断闪烁的错觉，我个人无法接受这种情况
• 虽然 gstatic.com 目前没有被墙（也有一说是直接被解析到国内服务器上了），但是我总觉得还是本地存储比较保险一些，然后再使用 CDN 来保证加载速度

要把 Google Fonts 中的第三方字体部署在本地也再容易不过了，只需要在 google-webfonts-helper 中选择想要的字体进行下载并存储在本地 FTP 指定目录下，然后把其提供的 CSS 代码添加到网页原有的样式后。以 WordPress 为例，可以把这段 CSS 代码复制到主题对应的 style.css 中，或者也可以新建一个 CSS 文件然后利用 wp_enqueue_style 或 wp_register_style 函数使得这个 CSS 文件可以被 WordPress 访问到。由于我使用的主题已经内置了 Google Fonts 的功能，所以我直接在 functions.php 中使用了一个同名函数覆盖掉了原有的从 fonts.googleapis.com/css 获取样式的函数：

function generate_google_fonts_link() {
    // replace remote stylesheet in fonts.googleapi.com with 
    // the local file
    return '/assets/google-fonts.css'; 
}

function load_google_font() {
if ( $fonts_link = generate_google_fonts_link() ) {
    wp_enqueue_style( 'google-fonts', $fonts_link, false );
}

add_action( 'wp_enqueue_scripts', 'load_google_font' );

其中的 /assets/google-fonts.css 就是 .woff 和 .woff2 字体文件所在的相对路径。

我的 google-fonts.css 的具体内容为：

/* josefin-sans-regular - latin_latin-ext_vietnamese */
@font-face {
  font-family: 'Josefin Sans';
  font-style: normal;
  font-weight: 400;
  src: local('Josefin Sans Regular'), local('JosefinSans-Regular'),
       url('/assets/fonts/josefin-sans-v14-latin_latin-ext_vietnamese-regular.woff2') format('woff2'), /* Chrome 26+, Opera 23+, Firefox 39+ */
       url('/assets/fonts/josefin-sans-v14-latin_latin-ext_vietnamese-regular.woff') format('woff'); /* Chrome 6+, Firefox 3.6+, IE 9+, Safari 5.1+ */
  font-display: swap;
}

/* noto-serif-sc-300 - chinese-simplified_latin */
@font-face {
  font-family: 'Noto Serif SC';
  font-style: normal;
  font-weight: 300;
  src: local('Noto Serif SC Light'), local('NotoSerifSC-Light'),
       url('/assets/fonts/noto-serif-sc-v6-chinese-simplified_latin-300.woff2') format('woff2'), /* Chrome 26+, Opera 23+, Firefox 39+ */
       url('/assets/fonts/noto-serif-sc-v6-chinese-simplified_latin-300.woff') format('woff'); /* Chrome 6+, Firefox 3.6+, IE 9+, Safari 5.1+ */
  font-display: swap;
}

/* noto-serif-sc-900 - chinese-simplified_latin */
@font-face {
  font-family: 'Noto Serif SC';
  font-style: normal;
  font-weight: 900;
  src: local('Noto Serif SC Black'), local('NotoSerifSC-Black'),
       url('/assets/fonts/noto-serif-sc-v6-chinese-simplified_latin-900.woff2') format('woff2'), /* Chrome 26+, Opera 23+, Firefox 39+ */
       url('/assets/fonts/noto-serif-sc-v6-chinese-simplified_latin-900.woff') format('woff'); /* Chrome 6+, Firefox 3.6+, IE 9+, Safari 5.1+ */
  font-display: swap;
}

font-display

如果你仔细一点看上面的 CSS 代码的话就会发现，不同于从 google-webfonts-helper 中直接复制出来的代码，我还为每个字体指定了一个 font-display 字段。font-display 是一个比较新的 CSS 属性，用来控制某一字体在尚未成功加载时采用何种方式显示文本。如果不明确指定 font-display，浏览器一般会采用「block」的方式进行显示，即，在该字体完全加载成功之前啥都不显示。

对于英文字体来说即使使用 block 问题也不大，因为哪怕是非本地的字体文件也就几十 KB，完全可以在几百甚至几十毫秒内完成加载，对用户几乎不会造成任何影响。但是对于中文字体来说这个问题就比较严重了（特别是对于我这种使用单个大体积 .woff2 文件的情况），因为在完成字体加载之前，通常会有一段肉眼可见的页面空白时间，对用户非常不用好：

为了改善这个情况，我们在使用第三方中文字体时，非常有必要把 font-display 设置为 swap 或者 fallback，前者会在当前字体文件加载时使用 font-family 中当前字体的后备字体进行文本显示，而后者则允许一个停顿时间（比如100ms），该停顿时间内先显示空白，如果过了这个时间阈值当前字体仍然没有完成加载，则采用和 swap 相同的策略进行文本显示。比如我们在 CSS 中对正文文本使用 font-family: 'Noto Serif SC', 'Microsoft Yahei', sans-serif，那么当思源宋体的字体文件还没加载完成时，浏览器会先用微软雅黑显示当前的正文文本（本机中找不到微软雅黑的话继续 fallback 到系统默认的非衬线字体），待思源宋体加载完毕后再立即切换过去：

除了 block、fallback 和 swap 之外，font-display 还有 optional 这个选项，即让浏览器自己做决定是否要对字体进行切换，对于一些网速比较慢的网页，第三方字体加载完成可能需要花费很长时间，这种情况下即使字体文件已经完成了加载浏览器也可能不再进行切换，而是继续以 fallback 字体进行显示。

最后总结一下我认为的目前可接受的中文 webfont 解决方案：

1. 字体文件在本地存储，用 CDN 进行加速
2. 合理设置 font-family，做好 fallback
3. 使用 font-display: swap 属性（optional 也可）

Reference

• Controlling Font Performance with font-display
• Everything you need to know about web fonts
• 思源宋体，如何评价，以及如何正确使用
• 分享 Web 中文字体应用指南
• Fonts.css — 跨平台中文字体解决方案
• Controlling Font Display Per Font-Face: the font-display descriptor