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:
<unknown>
shape in funcB:
<unknown>
shape in funcA:
(1, 2, 4)