关于python：TensorFlow的random_poisson只在CPU上运行


TensorFlow's random_poisson only runs on CPU

我正在尝试让 TensorFlow 的 random_poisson 函数在我的 GPU 上运行；鉴于这个 TensorFlow 源页面有一个函数 testCPUGPUMatch 可以比较在 CPU 和 GPU 上运行时 random_poisson 的输出，这似乎应该是可能的。但是，在使用代码进行测试时：

1
2
3
4
5
6
import tensorflow as tf

with tf.Session() as sess:
    with tf.device("/gpu:0"):
        test = sess.run(tf.random_poisson(1.0, [], dtype=tf.float64))
print(test)

我得到错误：

InvalidArgumentError: Cannot assign a device for operation
'random_poisson/RandomPoissonV2': Could not satisfy explicit device
specification '/device:GPU:0' because no supported kernel for GPU
devices is available. Registered kernels: device='CPU'; R in
[DT_INT64]; dtype in [DT_INT64] device='CPU'; R in [DT_INT64]; dtype
in [DT_INT32] device='CPU'; R in [DT_INT64]; dtype in [DT_DOUBLE]
device='CPU'; R in [DT_INT64]; dtype in [DT_FLOAT] device='CPU'; R
in [DT_INT64]; dtype in [DT_HALF] device='CPU'; R in [DT_INT32];
dtype in [DT_INT64] device='CPU'; R in [DT_INT32]; dtype in
[DT_INT32] device='CPU'; R in [DT_INT32]; dtype in [DT_DOUBLE]
device='CPU'; R in [DT_INT32]; dtype in [DT_FLOAT] device='CPU'; R
in [DT_INT32]; dtype in [DT_HALF] device='CPU'; R in [DT_DOUBLE];
dtype in [DT_INT64] device='CPU'; R in [DT_DOUBLE]; dtype in
[DT_INT32] device='CPU'; R in [DT_DOUBLE]; dtype in [DT_DOUBLE]
device='CPU'; R in [DT_DOUBLE]; dtype in [DT_FLOAT] device='CPU'; R
in [DT_DOUBLE]; dtype in [DT_HALF] device='CPU'; R in [DT_FLOAT];
dtype in [DT_INT64] device='CPU'; R in [DT_FLOAT]; dtype in
[DT_INT32] device='CPU'; R in [DT_FLOAT]; dtype in [DT_DOUBLE]
device='CPU'; R in [DT_FLOAT]; dtype in [DT_FLOAT] device='CPU'; R
in [DT_FLOAT]; dtype in [DT_HALF] device='CPU'; R in [DT_HALF];
dtype in [DT_INT64] device='CPU'; R in [DT_HALF]; dtype in
[DT_INT32] device='CPU'; R in [DT_HALF]; dtype in [DT_DOUBLE]
device='CPU'; R in [DT_HALF]; dtype in [DT_FLOAT] device='CPU'; R in
[DT_HALF]; dtype in [DT_HALF]

[[Node: random_poisson/RandomPoissonV2 =
RandomPoissonV2[R=DT_DOUBLE, S=DT_INT32, dtype=DT_DOUBLE, seed=0,
seed2=0, _device="/device:GPU:0"](random_poisson/shape,
random_poisson/RandomPoissonV2/rate)]]

没有列出已注册的 GPU 内核。当在我的 CPU 上运行时，代码的行为与预期一样，在我的 GPU 上运行时与 uniform_random 类似的代码也是如此。我是否以某种方式错过了 random_poisson 的 GPU 内核？是否不存在，即使链接的源页面暗示存在？如果不存在，是否有在 GPU 上运行的实现？这是目前我实现的一个相当复杂的模型中的瓶颈，所以修复它会很好。我在 Arch Linux 上的 Python 3.6.4 上运行 TensorFlow 1.8.0 版(从 pip 安装)，在 GeForce GTX 1050 上运行 CUDA 9.0 版和 cuDNN 7.0 版。

谢谢！

没有用于随机泊松的 GPU 内核。在 random_poisson_op.cc 中，只注册了 CPU 内核。

如果您查看链接的 testCPUGPUMatch 代码，它会调用 self.test_session(use_gpu=True, ...)，如果可能，它会尝试在 GPU 上运行操作。在后台，test_session 使用 allow_soft_placement 来执行此操作，如果无法在 GPU 上运行操作，则返回到 CPU。所以测试实际上是在 CPU 上运行 op。

[旁白：为什么我们有一个实际上什么都不做的测试？似乎 testCPUGPUMatch 是对许多不同随机操作的测试(参见 test_random_ops.py)，而 RandomPoisson 的实现者可能出于这个原因添加了它。似乎将其命名为 testCPUGPUMatchIfGPUImplementationExists...]

会更好

随时通过 github 问题提交 GPU 内核的功能请求，或提交带有实现的 PR。

相关讨论

1
2
3
4
5
6
7
import tensorflow as tf
config = tf.ConfigProto(allow_soft_placement=True)

with tf.Session(config=config) as sess:
    with tf.device("/gpu:0"):
        test = sess.run(tf.random_poisson(1.0, [], dtype=tf.float64))
print(test)
相关讨论