Tensorflow：在 Python 中编写一个 Op

Question

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

如何才能完整地写在Python中？

Answer 1

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

Wraps a python function and uses it as a tensorflow op.

Given a python function func, which takes numpy arrays as its inputs and returns numpy arrays as its outputs.

您的 python 函数需要：

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

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

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

Answer 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

Answer 3

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

Tensorflow：在 Python 中编写一个 Op

Tensorflow: Writing an Op in Python

tensorflow