Tensorflow：用 Python 编写操作答案

【问题标题】：Tensorflow: Writing an Op in PythonTensorflow：用 Python 编写操作
【发布时间】：2016-10-21 19:12:50
【问题描述】：

我想用 Python 写一个 Op。本教程仅说明如何使用 Python 包装器在 c++ 中完成此操作。 https://www.tensorflow.org/versions/master/how_tos/adding_an_op/index.html#adding-a-new-op

如何才能完全用 Python 编写？

【问题讨论】：

标签： tensorflow

【解决方案1】：

您可以使用tf.py_func(func, inp, Tout)。

包装一个python函数并将其用作张量流操作。

给定一个 python 函数 func，它将 numpy 数组作为其输入并返回 numpy 数组作为其输出。

你的python函数需要有：

numpy 数组作为输入，从带有参数inp 的图形馈送
numpy 数组作为输出，您需要在参数Tout 中将它们的类型指定给 TensorFlow

在函数内部，你可以做任何你想做的事情，if for 循环的条件，任何在 TensorFlow 中不可能的事情。

但是，该操作将在 CPU 上执行，因此它可能比 GPU 中等效的 TensorFlow 操作要慢。

【讨论】：

如何在GPU上执行操作？我必须在 C++ 上实现这个 GPU 版本吗？有没有办法用 Python 编写代码？
我认为没有办法在 Python 中编写 GPU 操作。您可以随时尝试使用您想到的实现提出另一个问题，并且有人可能会想出一个想法直接在 TensorFlow 中实现它。

【解决方案2】：

您可以使用tf.py_func 来调用python 函数。函数内部的操作也可以在 GPU 上。例如，我们可以纯粹在 python 中添加一个 Op 及其渐变，它在 GPU 上调用 Caffe：

def custom_loss_impl(x):
    caffe.set_mode_gpu()
    caffe.set_device(0)
    ...
    return np.float32(loss)

def custom_loss(x):
    tf.RegisterGradient("custom_loss_grad")(custom_loss_grad)
    g=tf.get_default_graph()
    with g.gradient_override_map({"PyFunc":"custom_loss_grad"}):
        return tf.py_func(custom_loss_impl,[x],[tf.float32])[0]

def custom_loss_grad_impl(x):
    caffe.set_mode_gpu()
    caffe.set_device(0)
    custom_loss_impl(x)
    ...
    return np.float32(gradient)

def custom_loss_grad(op,grad):
    x=op.inputs[0]
    return tf.py_func(custom_loss_grad_impl,[x],[tf.float32])#assume grad=1

【讨论】：

使用这种方法，并假设您的 tensorflow 图的其余部分在 GPU 上，您的数据将再次移动 tensorflow GPU -> CPU -> caffe GPU -> CPU -> tensorflow GPU.... .

【解决方案3】：

根据我的经验，不使用 C 编写新 Op 的主要原因是您想要实现自定义渐变。如果这就是您想要制作 Op 的原因，您现在可以使用 tf.custom_gradient 来实现

【讨论】：