如何保存tf 2.x变压器模型？

问题描述

简而言之，我遵循TF Transformer中的教程。

在此issue中保存模型时，其他人会发生此问题。

我的问题是1）如何编写正确的签名以保存模型？ 2）是否需要为模型@tf.function中的每个call函数添加class。

# tutorial: https://www.tensorflow.org/api_docs/python/tf/saved_model/save
infer_signature = transformer.call.get_concrete_function(
  tf.TensorSpec(shape=[None,None],dtype=tf.int64,name='encoder_input'),# encoder_input
  tf.TensorSpec(shape=[None,name='tar_input'),# tar_input
  tf.TensorSpec(shape=None,dtype=tf.bool,name='train'),# training
  tf.TensorSpec(shape=[4,dtype=tf.float32,name='enc_padding_mask'),# enc_padding_mask
  tf.TensorSpec(shape=[4,name='combined_mask'),# combined_mask
  tf.TensorSpec(shape=[4,name='dec_padding_mask')  # dec_padding_mask
)

saved_path = './saved_model'
transformer.save(saved_path,signatures=infer_signature)
# OR => tf.saved_model.save(transformer,saved_path,signatures=infer_signature)

错误：

_________________________________________________________________
Traceback (most recent call last):
  File "/Users/xiaofengwu/Google Drive/intern-zq/intern_notes/tvm_prj/tf_ocr_model_impl/transformer_tf_2/transformer_tf2.py",line 1038,in <module>
    transformer.save(saved_path,signatures=infer_signature)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py",line 1979,in save
    signatures,options)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py",line 134,in save_model
    signatures,options)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/save.py",line 80,in save
    save_lib.save(model,filepath,signatures,options)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py",line 976,in save
    obj,export_dir,options,Meta_graph_def)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/saved_model/save.py",line 1050,in _build_Meta_graph
    signature_serialization.canonicalize_signatures(signatures))
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_serialization.py",line 137,in canonicalize_signatures
    **tensor_spec_signature)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py",line 1073,in _get_concrete_function_garbage_collected
    self._initialize(args,kwargs,add_initializers_to=initializers)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py",line 697,in _initialize
    *args,**kwds))
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/eager/function.py",line 2855,in _get_concrete_function_internal_garbage_collected
    graph_function,_,_ = self._maybe_define_function(args,kwargs)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/eager/function.py",line 3213,in _maybe_define_function
    graph_function = self._create_graph_function(args,line 3075,in _create_graph_function
    capture_by_value=self._capture_by_value),File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py",line 986,in func_graph_from_py_func
    func_outputs = python_func(*func_args,**func_kwargs)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py",line 600,in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args,**kwds)
  File "/opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py",line 973,in wrapper
    raise e.ag_error_Metadata.to_exception(e)
ValueError: in user code:

    /opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_serialization.py:126 signature_wrapper  *
        structured_outputs,signature_function.name,signature_key)
    /opt/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_serialization.py:187 _normalize_outputs  **
        .format(original_outputs,function_name,signature_key))

    ValueError: Got non-flat outputs '(<tf.Tensor 'StatefulPartitionedCall:0' shape=(None,None,980) dtype=float32>,{'decoder_layer1_block1': <tf.Tensor 'StatefulPartitionedCall:1' shape=(None,8,4,None) dtype=float32>,'decoder_layer1_block2': <tf.Tensor 'StatefulPartitionedCall:2' shape=(None,'decoder_layer2_block1': <tf.Tensor 'StatefulPartitionedCall:3' shape=(None,'decoder_layer2_block2': <tf.Tensor 'StatefulPartitionedCall:4' shape=(None,'decoder_layer3_block1': <tf.Tensor 'StatefulPartitionedCall:5' shape=(None,'decoder_layer3_block2': <tf.Tensor 'StatefulPartitionedCall:6' shape=(None,'decoder_layer4_block1': <tf.Tensor 'StatefulPartitionedCall:7' shape=(None,'decoder_layer4_block2': <tf.Tensor 'StatefulPartitionedCall:8' shape=(None,None) dtype=float32>})' from 'b'__inference_call_54652'' for SavedModel signature 'serving_default'. Signatures have one Tensor per output,so to have predictable names Python functions used to generate these signatures should avoid outputting Tensors in nested structures.

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）

keras tensorflow transformer