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

如何解决如何保存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.

相关帖子：

Transformer model not able to be saved