支持 GPU 的 TensorFlow 操作

Question

有没有办法（或者可能是一个列表？）来确定所有提供 GPU 支持的 tensorflow 操作？ 现在，这对我来说是一个反复试验的过程——我尝试将一组操作放在 GPU 上。如果它有效，亲爱的。如果没有，请尝试其他地方。

这是迄今为止我发现的唯一相关（但没有帮助）：https://github.com/tensorflow/tensorflow/issues/2502

Answer 1

不幸的是，似乎没有内置方法可以获取确切的列表，甚至无法检查特定的操作。如上问题所述，调度是在原生 C++ 代码中完成的：如果相应的内核已注册到DEVICE_GPU，则可以将特定操作分配给 GPU。

我认为对您来说最简单的方法是 grep "REGISTER_KERNEL_BUILDER" -r tensorflow tensorflow 源库以获取匹配操作的列表，这看起来像 like this。

但请记住，即使使用 REGISTER_KERNEL_BUILDER 规范，也不能保证在 GPU 上执行操作。例如，无论现有内核如何，都会在 CPU 上分配 32 位 int Add。