TensorFlow 跨设备通信

Question

正如 TensorFlow 论文所述，Tensorflow 的跨设备通信是通过在设备中添加“接收节点”和“发送节点”来实现的。

据我了解，设备（请考虑仅涉及 CPU 设备）负责执行操作的计算。但是，数据（例如：操作产生的张量，变量缓冲区）驻留在内存中。我不知道如何物理地实现从一台设备到另一台设备的数据传输。我猜数据传输是通过共享内存实现的。那正确吗？

我将不胜感激有关如何实现数据传输的任何解释/相应代码。 PS：TensorFlow paper link，图4展示了跨设备通信机制。

Answer 1

在 TensorFlow 中，跨设备通信是使用 Rendezvous 接口实现的，该接口有多种不同的实现，具体取决于部署。该界面上的注释描述了总体思路：

// A Rendezvous is an abstraction for passing a Tensor
// from a producer to a consumer, where the consumer may safely
// request the Tensor before or after it has been produced.  A
// producer never blocks when using a Rendezvous.  A consumer has the
// choice of making a blocking call or providing a callback: in either
// case, the consumer receives the Tensor as soon as it is available.

正如您在问题中所指出的，TensorFlow 使用Send 和Recv 操作表示数据流图中的通信，这些操作会在图跨设备分区时自动添加到图中。对于在不同设备上具有源和目标的每条边，图分区器插入一对 Send 和 Recv 操作，它们共享相同的“集合键”（自动生成的字符串名称，用作集合点的待传达张量的索引）。 implementation of the Send op 很简单：它调用Rendezvous::Send()，传入它的集合键和单个输入张量，然后立即返回而不阻塞。 implementation of the Recv op 稍微复杂一些：它注册一个回调，当具有给定键的张量可用时调用。 callback 负责“生成”Recv 操作的输出，并解除对后续计算的阻塞。

Rendezvous 实现执行传输数据的实际工作：

• IntraProcessRendezvous 在同一进程中处理设备之间的数据传输。如果（不太可能）在同一进程中的两个 CPU 设备之间进行传输，则可以通过simple Tensor assignment 实现传输。否则，TensorFlow kicks off 一个特定于设备的 DMA 例程，用于在 CPU 和 GPU 设备之间传输数据。

• BaseRemoteRendezvous 类及其子类在发送方和接收方可以在不同进程中的情况下处理跨设备通信。这个类的主要实现是RpcRemoteRendezvous，其中uses gRPC来处理远程传输。