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

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

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

你的python函数需要有:

• numpy数组作为输入,从带有参数的图形中提供 inp
• numpy数组作为输出,你需要在参数中指定它们的类型为TensorFlow 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