如何知道在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>);

这意味着GPU上softmax的内核类型为float16。前提条件是您必须在启用GPU的情况下构建张量流。