tensorflow - 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函数并将其用作tensorflow op。
给定一个python函数func，该函数将numpy数组作为其输入并返回numpy数组作为其输出。


您的python函数需要具有：

numpy数组作为输入，由带有参数inp的图提供
numpy数组作为输出，您需要在参数Tout中将其类型指定为TensorFlow

在函数内部，如果条件为for循环，则可以执行任何您想做的事，这在TensorFlow中是不可能完成的。

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