为什么 tf2 不能将 tf_function 模型保存为 .pb 文件？

Question

我试图在 official website 上保存一个类似 transformer 官方代码的模型，但是当我想保存 train_step 图形或使用 tf.summary.trace_on 在其上进行跟踪时，它 errors.the 错误为

    ---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py in convert(x)
    936         try:
--> 937           x = ops.convert_to_tensor_or_composite(x)
    938         except (ValueError, TypeError):

15 frames
TypeError: Can't convert Operation 'PartitionedFunctionCall' to Tensor (target dtype=None, name=None, as_ref=False)

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py in convert(x)
    941               "must return zero or more Tensors; in compilation of %s, found "
    942               "return value of type %s, which is not a Tensor." %
--> 943               (str(python_func), type(x)))
    944       if add_control_dependencies:
    945         x = deps_ctx.mark_as_return(x)

TypeError: To be compatible with tf.contrib.eager.defun, Python functions must return zero or more Tensors; in compilation of <function canonicalize_signatures.<locals>.signature_wrapper at 0x7fcf794b47b8>, found return value of type <class 'tensorflow.python.framework.ops.Operation'>, which is not a Tensor.

我以为是张量运算出错，写了一个demo来证实我的想法：

    import tensorflow as tf
sig=[tf.TensorSpec(shape=(None, None), dtype=tf.int64),tf.TensorSpec(shape=(None, None), dtype=tf.int64)]
@tf.function(input_signature=sig)
def cal(a,d):
  b=a[1:]
root=tf.Module()
root.func = cal

# 获取具体函数。
concrete_func = root.func.get_concrete_function(
      tf.TensorSpec(shape=(None, None), dtype=tf.int64),tf.TensorSpec(shape=(None, None), dtype=tf.int64)
)
tf.saved_model.save(root, '/correct', concrete_func)

错误如期发生。但是我该如何解决呢？ positional_encoding 需要，但我不知道如何替换此操作。

Answer 1

我猜是因为我没有在pb文件中保存transformer模型

为什么 tf2 不能将 tf_function 模型保存为 .pb 文件？

Why tf2 can't save a tf_function model as .pb file?

graph

tensorflow2.0