首页 > 解决方案 > 为什么 tf2 不能将 tf_function 模型保存为 .pb 文件?

问题描述

我试图在官方网站上保存一个类似于transformer的官方代码的模型,但是当我想保存train_step图或使用tf.summary.trace_on对其进行跟踪时,它会出错。错误为

    ---------------------------------------------------------------------------
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.

我认为这是张量操作的一些错误,并编写了另一个演示来确认我的想法:

    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 需要,我不知道如何替换此操作。

标签: graphtensorflow2.0

解决方案


我想它是因为我没有将变压器模型保存在 pb 文件中


推荐阅读