tensorflow - 没有用于 int32 变量操作的 GPU 内核

Question

当我用整数张量初始化Variable操作时，它似乎没有GPU内核。例如，运行:

with tf.device('/gpu:0'):
  var = tf.Variable(initial_value=([[1,2],[3,4]]))
sess = tf.Session()
sess.run(var)

会导致异常:

tensorflow.python.framework.errors.InvalidArgumentError: Cannot assign a device to node 'Variable': Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available

这是否意味着 TensorFlow 中的所有变量都存储在与 CPU 设备关联的内存(我猜是 RAM)上，而不是 GPU 内存上？当我们只使用一个 GPU 来训练模型时，是不是会更慢，因为变量值必须从 RAM 复制到 GPU 内存多次？

谢谢!

Answer 1

TL;DR: GPU 上没有用于 int32 值变量运算的内核，但有其他元素类型 (float16/float32/float64)在 GPU 上受支持。

---

TensorFlow 目前仅支持 (half-, single-, and double-precision)浮点 variables on GPU 。事实上，总的来说，TensorFlow 对 GPU 上的 int32 操作的支持有限，您会看到像这样的有趣评论 throughout the C++ code :

// A special GPU kernel for int32.
// TODO(b/25387198): Also enable int32 in device memory. This kernel
// registration requires all int32 inputs and outputs to be in host memory.

在 GPU 上对 int32 数据运行操作并不是特别困难，那么为什么 TensorFlow 不支持它呢？

1. TensorFlow 中的默认放置算法会贪婪地将操作放置在该操作的“最佳可能”设备上。这意味着如果 GPU 可用，则在 GPU 上，否则在 CPU 上。
2. 根据经验，我们发现真实模型中的大多数 int32 操作都会生成和消耗“元数据”(如张量的形状)，并且此信息通常位于主机内存中。这些操作出现在梯度代码中，以处理可变大小的张量。
3. 在 GPU 上注册 int32 内核意味着必须将这些小型元数据张量复制到设备内存或从设备内存复制，以便运行通常非常便宜的操作。此外，将结果复制回主机内存可能会停止 GPU 流，从而导致性能大幅下降。

因此，我们做出了妥协，删除了许多 int32 内核的注册，同时我们开发了一种更智能的放置算法，可以识别操作何时无法从在 GPU 上运行中获益。由于在 GPU 上为 int32 注册的操作非常少，因此将变量放在 GPU 上几乎没有任何好处，会消耗相对稀缺的内存。