如何知道tensorflow中哪些操作不能放在GPU上？答案

【问题标题】：How can I know which operation can not be placed on GPU in tensorflow?如何知道tensorflow中哪些操作不能放在GPU上？
【发布时间】：2016-11-19 23:53:06
【问题描述】：

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

谢谢

【问题讨论】：

标签： tensorflow gpu

【解决方案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>);

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

【讨论】：

非常感谢！这真的很有帮助！
在某些内核中，我看到了类似#define REGISTER_GPU_KERNELS(type) 的内容。我如何知道启用了哪些类型？
在其下，有一个宏TF_CALL_GPU_NUMBER_TYPES，在repo中搜索，你会找到this page。它表示TF_CALL_GPU_NUMBER_TYPES 将为float16、float32 和double 类型注册内核。