tf.py_function使用心得

tensorflow为在tensor环境中插入python代码，提供了几个function，如下：

1.tf.py_func：tf1下的api，用于在静态图中嵌入python代码。

2.tf.py_function：tf2下的api，用于在tf2环境中嵌入python代码。

3.tf.numpy_function：tf2下的api，用于在tf2环境中嵌入numpy代码。

以上三个函数的使用方法和限制条件可以参考tf官方doc。

然而实际使用时笔者发现限制比官方doc描述的要多，在autograph模式下，调用上面的方法大多数情况下是不起作用的。举例如下：

当前code在eager环境中，有如下调用关系：funcA call funcB call ... call funcN, 其中funcN中调用了tf.py_function， 如果想要这些函数都在autograph下运行，那么使用tf.function装饰funcA，则funcN中的tf.py_function将失效，只有用tf.function装饰funcN，才有效。也就是说只有funcN能在aotugraph下运行。另外，如果使用tf.function装饰funcA，同时又用tf.function装饰funcN，则funcN中的tf.py_function也将失效。

通过进一测试分析，笔者发现，上面的不起作用或者失效其实是假象，实际上tf.py_function有被执行，只不过其返回的是eager tensor object，由于在autograph环境中，eager tensor是不可见的（不知道shape，无法访问其内容），如果这个eager tensor跟当前autograph环境中的tensor进行计算，可能会报错。如果将eager tensor返回给非autograph环境，择该eager tensor又变得可见了。举例如下：

当前code在eager环境中，有如下调用关系：funcA call funcB call ... call funcN, 其中funcN中调用了tf.py_function， 如果使用tf.function装饰funcB，则funcA为eager环境，funcB到funcN均为autograph环境。此时funcN中tf.py_function的返回值为eager tensor，在funcB到funcN中均不可见，在funcA中可见。sample code如下：

output:

shape in funcC:

shape in funcB:

shape in funcA:

(1, 2, 4)