TensorFlow 设备上下文

Question

我正在向 TensorFlow 添加新硬件，但我遇到了设备上下文问题。

只是为了确保我理解它们的目的，从代码中，它们似乎在管理计算图中给定节点的设备和数据。也就是说，输入数据和数据输出的设备在设备上下文中指定，执行器处理这些设备之间来回传递的张量。每个 OpKernelContext 似乎都包含一个设备上下文，用于控制给定内核的特定执行。

目前对于我的新硬件，设备上下文都是空的，这最终会导致段错误。根据现有 C++ 的风格，我希望有允许我为特定设备类型（即 GPU、CPU）“注册”设备上下文的宏，但我找不到这些。那么我的问题是，当它们是为我的设备创建时，如何将正确的设备上下文添加到 OpKernelContext 中。

请注意，我没有编写特定于我的硬件的设备上下文类。我注意到 ThreadPoolDevice 似乎没有特定于它的 DeviceContext 实现。我认为这是因为基类 DeviceContext 是为 ThreadPools 实现的。

如果对 DeviceContexts 进行任何澄清，我将不胜感激。

Answer 1

DeviceContext 对象 有两个目的：

1. 目前，大多数调用者可能会忽略某些 StreamExecutor 特定的位（gpu::Stream / MaintainLifetime），因为它们特定于 GPU。

   您需要在 OpKernels 中处理底层设备资源，并且 DeviceContext 对象包含用于计算的“流”对象。

   我们尚未在 DeviceContext 中实现其他设备应作为其资源实现的不透明句柄，但这正是我们所需要的。因此，不幸的是，在我们可以让非基于 StreamExecutor 的设备工作之前，这是一个 TODO。

2. 另一个组件是处理从设备复制到 CPU 以及从设备复制到 CPU 的代码。您说得对，CPU 不存在 DeviceContexts，因为 CPU 是主机设备，从这个角度来看，不需要特殊对待。

   但作为一个例子，我们可以看看 GPU 的代码。 tensorflow/core/common_runtime/gpu_device_context.h 是 GPU 设备上下文的一个示例，它实现了 DeviceContext 接口。接口的实现是here，它委托给GPUUtil 类中的代码来实际执行内存拷贝。它碰巧使用 StreamExecutor 框架来处理底层副本，但您自己的设备将使用您拥有的任何 API 来复制到主机和从主机到设备。

   如果你想从一个设备复制到另一个设备，有一个特殊的注册。 GPU 到 GPU 的示例是： here 用于实现，here 用于注册该功能。

在某些时候，我们可能会将其重构为更简洁，因此注册都是统一的（CPU 到设备、设备到 CPU、设备到设备）。目前它有点临时。

因此，目前支持其他设备的工作正在进行中，但我们很高兴与您和其他人合作以充实这种支持。