3.3 身份图重写：改善搜索空间以提高峰值内存占用
重新组织不规则连接神经网络的计算图可能会显著减少计算过程中的峰值内存占用µpeak。例如，值得注意的是，大量基于 NAS 的研究（Liu et al., 2019a; Zhang et al., 2019）大量使用连接作为合并来自输入激活的多个分支信息的自然方法，从而扩大神经架构的搜索空间。然而，具有许多入边的连接可能会延长输入激活的存活时间并增加内存压力，这在资源受限的场景中尤其不利。为了解决这个问题，我们提出身份图重写，以有效地减少连接周围的µpeak，同时保持算术输出不变。为此，我们提出了两个主要示例，展示在不规则连接神经网络中受益于我们技术的图模式：
通道划分（卷积）
在不规则连接神经网络中，一个典型模式是连接（concat: [·]），它在卷积（conv: ∗）之前从多个输入分支中获取输入。在执行这种模式时，峰值内存占用 ( \mu_{\text{peak}} ) 发生在计算输出 y ∈ R^n 的同时，连接的输入x∈R ^n 也必须驻留在内存中。我们的目标是实现与连接相同的算术结果和逻辑效果，但规避相应的看似过高的内存成本。为此，我们对跟随连接的卷积进行通道划分，以便分割的卷积可以在输入 x_i 可用时立即计算。
方程 3-6 详细描述了这种替换的数学推导。具体而言，如方程 3 所示，每个内核迭代并汇总在卷积中对通道的卷积结果。然而，利用乘法的分配性质 Pi 和 * ，这些可以转化为通道划分卷积的求和，我们称之为部分卷积（partial conv）。这种部分卷积从图中移除了连接，从而降低了内存成本。如图 9 所示，相同计算的内存成本从 Pxi+y 降低到 max(w_i∗x_i)+y ，当连接的入边数量增多时，这种效果更加显著。
内核划分（深度卷积）
深度卷积（depthwise convolution）(Sifre & Mallat, 2014; Howard et al., 2017) 已被证明在减少计算的同时实现竞争性能，因此在以极高效率为主要目标的网络中广泛使用。对于连接 （concat）后跟深度卷积（depthconv），与上面的连接+卷积情况类似，峰值内存占用 µpeak 发生在连接的 x 在内存中，而结果 y 在 x 被卸载之前也被保存到内存中。这次，我们利用不同内核之间的独立性，对跟随连接的深度卷积进行内核划分，以便每个输入 x_i 被计算为较小的特征图，而不在内存中驻留太久。因此，方程 7-8 推导了这种替换。方程 7 显示 y 中的每个组件都是独立的（不同的下标索引），且适合进行划分。换句话说，这种重写简单地揭示了深度卷积与连接之间的交换性质，加上内核划分，从而显著降低 µpeak。
实现
遵循编译器中使用模式匹配算法的一般实践（Lattner & Adve, 2004; Rotem et al., 2018; Jia et al., 2019），我们使用模式匹配实现身份图重写，以识别可以替换为计算成本更低的操作的图区域。同样，我们利用这一技术识别导致较低内存成本的区域。
4 评估
我们用四个代表性的不规则连接神经网络图评估 SERENITY。我们首先比较 SERENITY 与 TensorFlow Lite（Google）在使用相同线性内存分配方案下的峰值内存占用。此外，我们还实验了峰值内存占用减少对外部内存通信的影响。我们还对使用 SwiftNet Cell A（Zhang et al., 2019）所提议的基于动态编程的调度器和图重写的收益进行了深入分析。最后，我们研究了自适应软预算对调度时间的影响。
4.1 方法论
基准和数据集。表 1 列出了用于评估的网络的详细信息，这些网络是来自神经架构搜索（NAS）和随机网络生成器（RAND）中不规则连接神经网络的代表：DARTS（Liu et al., 2019a）用于 ImageNet，SwiftNet（Zhang et al., 2019）用于包含人类存在或缺失（HPD）的数据集，RandWire（Xie et al., 2019）用于 CIFAR10 和 CIFAR100。DARTS（Liu et al., 2019a）是一种基于梯度的 NAS 算法。特别是，我们专注于用于 ImageNet 图像分类的学习正常单元：仅第一个单元，因为它具有最高的峰值内存占用，网络的其余部分只是重复堆叠相同的单元，遵循 NASNet 的实践（Zoph et al., 2018）。SwiftNet（Zhang et al., 2019）是 NAS 生成的网络，旨在针对人类检测数据集。RandWire（Xie et al., 2019）来自随机网络生成器，用于在 CIFAR10 和 CIFAR100 上进行图像分类。该表还列出了它们的数据集、乘法累加计数（# MAC）、参数数量（# WEIGHT）和在各自数据集上的 top-1 准确率。
4.2 实验结果
与 TensorFlow Lite 的比较。图 10 评估了 SERENITY 相对于 TensorFlow Lite 在上述网络的不同单元上的内存占用减少情况。这些图表明，SERENITY 的基于动态编程的调度器将内存占用减少了 1.68 倍，而对图没有任何更改。此外，所提议的图重写技术在峰值内存占用方面平均减少了 1.86 倍（额外减少 10.7%）。结果表明，SERENITY 在不规则连接神经网络的峰值内存占用方面实现了显著降低。
图 10. SERENITY 相对于 TensorFlow Lite 的峰值内存占用减少（无内存层次）。
图 11. SERENITY 相对于 TensorFlow Lite 的外部内存通信减少（有内存层次）。
外部内存通信的改善。我们还展示了 SERENITY 如何影响外部内存通信，这在很大程度上影响了功耗和推理速度（Chen et al., 2016; Gao et al., 2017; Sharma et al., 2018）。为此，图 11 扫描了不同的片上内存配置，以测量在具有多级内存层次的系统上外部通信的减少。由于我们事先知道整个调度，我们使用 Belady 的最优算法（Belady, 1966），也称为先知算法，测量外部内存通信，从而提取提议调度的效应。结果显示，SERENITY 可以在具有 256KB 片上内存的设备上将外部内存通信减少 1.76 倍。特别是，尽管有一些情况下峰值内存占用已经足够小以适合片上内存（在图中标记为 N/A），但也有一些情况，SERENITY 通过成功将激活内容限制在片上内存中消除了外部通信，而 TensorFlow Lite 则未能做到这一点（在图中标记）。这表明，SERENITY 在减少内存占用方面的努力在具有内存层次的系统中同样有效地降低了外部内存通信，从而降低了功耗和推理速度。
(a) 带有内存分配器的内存占用（TensorFlow Lite 的峰值内存占用 = 551.0KB）。
(b) 无内存分配器的内存占用。
图 12. 在使用和未使用内存分配器的情况下运行 SwiftNet Cell A 时的内存占用（红色箭头表示减少）。
来自基于动态编程的调度器和身份图重写的改善。为了展示改进的来源，图 12 绘制了运行 SwiftNet Cell A 时的内存占用。图 12(a) 显示了 SERENITY 在内存分配下的内存占用。该图显示，SERENITY 的基于动态编程的调度器对峰值内存占用带来了显著改善（551.0KB→250.9KB），而图重写通过利用减少大内存占用区域的模式进一步改善了这一结果（250.9KB→225.8KB）。为了关注调度器和图重写的影响，图 12(b) 展示了 SERENITY 在没有内存分配情况下的内存占用：运行网络时激活的总和。该图显示，所提议的调度器在不更改图的情况下找到了具有最优（最小）峰值内存占用的调度。接着，它显示，所提议的图重写可以进一步将峰值内存占用减少 12.5KB（200.7KB→188.2KB）。结果表明，改进的显著部分来自于所提议的基于动态编程的调度器和图重写。