tf.py_func的一些使用笔记——TensorFlow1.x

tensorflow.py_func是TensorFlow1.x版本下的函数，在TensorFlow.2.x已经不建议使用了，但是依然可以通过tf.compat.v1.py_func的方式来进行调用。

可以说TensorFlow1.x下的py_func函数在TensorFlow2.x下除了通过tf.compat.v1.py_func的方式来进行调用就再也没有等价的使用方法了，具体可以看TensorFlow2.x的API文档：

https://tensorflow.google.cn/api_docs/python/tf/compat/v1/py_func

----------------------------------------------------------

这里需要着重说明一点，很多人认为TensorFlow1.x中的tf.py_func等价于TensorFlow2.x中的tf.py_function，其实不然。在TensorFlow2.x中除了对tf.py_func进行v1版本保留和兼容的tf.compat.v1.py_func，其实是没有完全同tf.py_func等价的函数。如果说在 TensorFlow2.x中比较相近的函数应该是tf.numpy_function而不是tf.py_function。

在TensorFlow2.x中对tf.numpy_function的解释：

https://tensorflow.google.cn/api_docs/python/tf/numpy_function

在TensorFlow2.x中对tf.py_function的解释：

https://tensorflow.google.cn/api_docs/python/tf/py_function

----------------------------------------------------------

TensorFlow2.x中 tf.numpy_function 和 TensorFlow1.x中 tf.py_func（tf.compat.v1.py_func）中唯一的区别是：

tf.py_func中是可以设置函数是否考虑状态的，而tf.numpy_function中是必须要考虑状态的（没有定义不考虑状态的设置）。

This name was deprecated and removed in TF2, but tf.numpy_function is a near-exact replacement, just drop the stateful argument (all tf.numpy_function calls are considered stateful).

=============================================

对TensorFlow1.x中 tf.py_func进行下一步解释：

tf.py_func其实是将python函数包装成TensorFlow的一个操作operation，tf.py_func的输入可以是numpy，可以是tensor，也可以是Variable，其输入只能是tensor。

tf.py_func定义的操作是属于TensorFlow的计算图的，在定义tf.py_func时是不会具体执行的，只有在具体的tf.Session中还可以执行，但是tf.py_func并不同于其他的TensorFlow的operation，因为tf.py_func定义的操作是运行在python空间下的而不是运行在TensorFlow空间下的。

tf.py_func定义后，在session中运行时的基本原理就是将输入的变量（不论是tensor还是numpy.array）转换为python空间下的numpy.array变量，在经过numpy运算后在将获得的numpy.array结果转换为tensor，给到TensorFlow的计算图。

其实，tf.py_func的功能完全可以手动实现类似的，就是手动的把tensor变量转为numpy.array，然后运算好后把结果手动转为tensor，tf.py_func最大的好处就是把这一过程给自动化了，不过随之也使这个运算过程变得难以理解了。从tf.py_func的原理我们就可以知道，虽然tf.py_func可以作为TensorFlow计算图的一部分挂在计算图上，但是由于其本质是将TensorFlow空间变量转为python空间变量后经过运算再转为TensorFlow空间变量，中间经过了命名空间和运算空间的转换，因此tf.py_func是不可以进行梯度反传的，或许我们更可以把这个操作看做是一种简易的为TensorFlow提供支持的python库。

================================================

2022年10月13日更新

如果tf.py_func包装的python函数的参数是string类型，那么传到包装的函数内时会被自动转为bytes类型，也就是string变bytes，这一点需要注意，否则真的是不知道什么地方报错的。

例子：

import tensorflow as tf
import numpy as np

sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    print(a, b)
    return np.array(len(a+b), dtype=np.float32)

x = "abc"
y = "bde"
ans = tf.py_func(fun, (x, y), (tf.float32, ), name="ab_op")
print("="*30, "result:")
print(ans)

print(sess.run(ans))

运行结果：

可以看到，由tensorflow空间传参到python空间会自动的将string类型转为bytes类型。

在python3.x版本中，可以使用bytes.decode()的方法将传入的bytes类型转会string类型，具体：

修改后的代码：

import tensorflow as tf
import numpy as np

sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    print(a, b)
    a = a.decode()
    b = b.decode()
    print(a, b)
    return np.array(len(a+b), dtype=np.float32)

x = "abc"
y = "bde"
ans = tf.py_func(fun, (x, y), (tf.float32, ), name="ab_op")
print("="*30, "result:")
print(ans)

print(sess.run(ans))

View Code

重点部分：

================================================

一些例子：

以下代码均为TensorFlow1.x版本：

import tensorflow as tf
import numpy as np

sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    return a+1, b+1

x = np.array([1.0,2.0,3.0], dtype=np.float32)
y = np.array([4.0,5.0,6.0], dtype=np.float32)

ans=tf.py_func(fun, (x, y), (tf.float32, tf.float32), name="ab_python")

print("="*30, "result:")
print(ans)
print(sess.run(ans))

运行结果：

可以看到，tf.py_func的执行其实是为TensorFlow定义了一个operation，而tf.py_func所包装的python函数只有在TensorFlow执行计算图的时候才会被真正执行。

tf.py_func为包装的python函数所传入的参数可以是numpy.array类型，也可以是tensor类型，也可以是Variable类型，但是不管在tf.py_func中传入的参数是什么类型，最后传入到所包装的python函数中都会被转为numpy.array类型，而包装后的函数在session开始执行后所返回给计算图的数据类型也会被转换为tensor类型。

--------------------------------------------------------

包装的参数为tensor:

import tensorflow as tf
import numpy as np


sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    return a+1, b+1

x = tf.constant([1.0,2.0,3.0], dtype=np.float32)
y = tf.constant([4.0,5.0,6.0], dtype=np.float32)
ans =tf.py_func(fun, (x, y), (tf.float32, tf.float32), name="ab_op")

print("="*30, "result:")
print(ans)
print("session is running!!!")
print(sess.run(ans))

运行结果：

--------------------------------------------------------

包装的参数为Variable:

import tensorflow as tf
import numpy as np

sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    return a+1, b+1

x = tf.Variable([1.0,2.0,3.0], dtype=np.float32)
y = tf.Variable([4.0,5.0,6.0], dtype=np.float32)
ans = tf.py_func(fun, (x, y), (tf.float32, tf.float32), name="ab_op")
print("="*30, "result:")
print(ans)

sess.run(tf.global_variables_initializer())
print(sess.run(ans))

运行结果：

---------------------------

包装的参数为Variable:

import tensorflow as tf
import numpy as np

sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    return a+b

x = tf.Variable([1.0,2.0,3.0], dtype=np.float32)
y = tf.Variable([4.0,5.0,6.0], dtype=np.float32)
ans = tf.py_func(fun, (x, y), (tf.float32, ), name="ab_op")
print("="*30, "result:")
print(ans)

sess.run(tf.global_variables_initializer())
print(sess.run(ans))

运行结果：

---------------------------------------------

包装的参数为Variable，求反传梯度报错:

import tensorflow as tf
import numpy as np

sess = tf.Session()

def fun(a, b):
    print("+"*30)
    print("function fun excute!!!")
    return a+b

x = tf.Variable([1.0,2.0,3.0], dtype=np.float32)
y = tf.Variable([4.0,5.0,6.0], dtype=np.float32)
x2 = tf.Variable([1.0,2.0,3.0], dtype=np.float32)
y2 = tf.Variable([4.0,5.0,6.0], dtype=np.float32)
ans = tf.py_func(fun, (x, y), (tf.float32, ), name="ab_op")
ans2 = x2 + y2
print("="*30, "result:")
print(ans)

sess.run(tf.global_variables_initializer())
print(sess.run(ans))

op2 = tf.gradients(ans2, (x2, y2))
print("Ops2 Gradients: \n", sess.run(op2))

op = tf.gradients(ans, (x, y))
print("Ops Gradients: \n", sess.run(op))