现代机器学习框架支持在多种硬件设备上执行运行，而不同设备间的内存往往是独立管理的，即在某个时刻数据可能存储在不同的设备内存中。那么当网络的不同部分在不同的硬件设备上运行时，各个设备不一定总是拥有最新的数据，因而存在设备间同步数据的问题。本文分析各个框架的内存同步机制。

直观地，设备间的内存同步机制应当达到两个目标：

1. 尽可能减少内存拷贝。即如无必要，则不进行任何同步。
2. 尽量减少资源分配。即如无必要，用于描述同一数据的内存不在多个设备上分配。

运行在机器学习框架之上的神经网络是一张有向图，其中的节点是计算操作、边是操作间的张量。当这些节点在不同的设备上运行时，除了“边界”张量所描述的数据可以在不同的设备上有拷贝，其他的张量所描述的数据应当都只有一份拷贝。所谓“边界”张量是指连接两个分别在不同硬件设备上计算的张量。

Caffe 中的 SyncedMemory

SyncedMemory 是 Caffe 中专门负责内存同步的模块。SyncedMemory本质上是个状态机，包含四种状态：

• UNINITIALIZED：未初始化。尚未分配任何资源
• HEAD_AT_CPU：CPU 拥有最新的数据
• HEAD_AT_GPU：GPU 拥有最新的数据
• SYNCED：CPU 和 GPU 都拥有最新的数据

下面是状态转换图。

SyncedMemory 同时减少了资源分配和内存拷贝。当张量不是“边界”时，它们的状态转换只有一种。例如当张量在 CPU 端时，只有从 UNINITIALIZED 转换到 HEAD_AT_CPU。这样不需要在 GPU 端分配资源，也无需拷贝数据。数据拷贝只有在从 HEAD_AT_CPU 或 HEAD_AT_GPU 转换到 SYNCED 时才会发生。

Caffe 的 SyncedMemory 看起来简单直观，但在集成更多的硬件设备时会遇到一些麻烦。例如，当我们需要在三种硬件设备间同步数据时，SYNCED 状态的语义就不够明确。一般而言，同步数据的过程是将数据从一个设备拷贝到另一个设备。在这一拷贝过程结束后，数据可能只在两个设备间是同步的，因此我们可能需要引入更多的表示同步的状态。例如，当引入 NPU 之后，我们需要引入三种额外的同步状态：SYNCED_CPU_NPU、SYNCED_CPU_GPU、SYNCED_GPU_NPU。状态转换图如下。

上图的转换只是理想中的情况，实际上 GPU 和 NPU 之间可能无法直接同步数据。因为 GPU 和 NPU 这类相关底层库一般只会提供与 CPU 交互的拷贝接口，那么在实践中可能所有的数据拷贝 CPU 都会参与。即从 GPU 到 NPU 的拷贝分两步，首先将数据从 GPU 拷贝到 CPU，然后再从 CPU 拷贝到 NPU。那么在GPU、CPU、NPU 的例子中，便不存在 SYNCED_GPU_NPU 状态，因为 CPU 端的数据也是同步的。这种情况的转换图如下所示。

类似地，当增加更多的设备时，需要对 SyncedMemory 进行大量的改造，添加足够多的同步“中间”状态。

在这些状态转换中，只有当第一次到达 HEAD_AT_[CPU|GPU|NPU]时，才会在相应的设备上分配资源。这些资源一旦分配，一般会在 SyncedMemory 对象的生命周期中留存。