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)。

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

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

你的python函数需要：

numpy数组作为输入，使用参数inp
numpy数组作为输出，您需要在参数Tout

在函数内部，你可以做任何你喜欢的事情，如果for循环的条件，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 的主要原因是要实现自定义渐变。如果这就是为什么要进行手术，那么现在可以使用tf.custom_gradient

Tensorflow：用Python编写Op

3 个答案: