Autograph 无法转换 if 导入笔记本中的语句

Autograph fails to convert if statements from imported notebook

在笔记本中 FuncDef.ipynb 我有以下代码:

import tensorflow as tf

@tf.function
def myfunc(x):      
    if x < 0:
        x = -x
    return x

x = tf.constant(-1)
print(myfunc(x))

如果我执行它,它会按预期运行并打印 tf.Tensor(1, shape=(), dtype=int32).

现在我尝试通过将它导入另一个笔记本来执行相同的代码:

from google.colab import drive
drive.mount('/content/drive')
%cd "/content/drive/My Drive/Colab Notebooks"

!pip install import-ipynb
import import_ipynb
import FuncDef

失败并出现以下错误:

OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

从导入的代码段执行时,Autograph 似乎无法转换 if 语句。有解决办法吗?当然有摆脱 if 语句的简单解决方案,但我不想这样做,因为它会使我的实际代码更难阅读。

问题很可能出在 import-ipynb 模块中。

证明:

文件x.py

import tensorflow as tf

@tf.function
def myfunc(x):
    if x < 0:
        x = -x
    return x


x = tf.constant(-1)
print(myfunc(x))

文件y.py

import tensorflow as tf

from x import myfunc

x = tf.constant(-1)
print(myfunc(x))

python y.py的输出是:

tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(1, shape=(), dtype=int32)

符合预期。

问题可能是 import-ipynb 如何处理装饰器? (这只是一个猜测)。

looks like单元格的代码执行就是这两行

                # transform the input to executable Python
                code = self.shell.input_transformer_manager.transform_cell(cell.source)
                # run the code in themodule
                exec(code, mod.__dict__)

但我不明白为什么当有装饰器时这段代码会失败(但也许 tf.function 装饰器正在某处创建一些全局对象,这种加载和执行代码的方式不在设计范围内tf.function?)

如果您可以尝试不使用笔记本(只需将笔记本代码导出到 .py 文件中),您就可以解决这个问题。