我怎么知道tensorflow中哪些操作不能放在GPU上？

Question

在tensorflow中如何知道哪些操作不能放在GPU上？有我可以检查的地方吗？

谢谢

Answer 1

您可以检查位于此目录的操作的内核（即设备上的实现）：https://github.com/tensorflow/tensorflow/tree/r0.11/tensorflow/core/kernels/

例如，假设您想知道 softmax 是否可以放在 GPU 上。您可以导航到 softmax 的内核：https://github.com/tensorflow/tensorflow/blob/r0.11/tensorflow/core/kernels/softmax_op.cc。您会发现以下代码：

REGISTER_KERNEL_BUILDER(
    Name("Softmax").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"),
    SoftmaxOp<GPUDevice, Eigen::half>);

这意味着 softmax 在 GPU 上有一个 float16 类型的内核。先决条件是您必须在启用 GPU 的情况下构建您的 tensorflow。

How can I know which operation can not be placed on GPU in tensorflow?