Fail to save Bert2Bert model instance without passing the label input i.e. target_ids

LoicDagnas commented 3 years ago

1. The entire URL of the file you are using

https://github.com/tensorflow/models/blob/master/official/nlp/nhnet/models.py

2. Describe the bug

When saving a Bert2Bert model instance, I am forced to pass the target_ids input which should be used only at training time.

Passing the target ids calling the model and using a custom serving function without these target ids as input make the saving works but IMO it is weird to do it that way.

3. Steps to reproduce

You can simply run the following piece of code:

import tensorflow as tf
from official.nlp.nhnet.configs import UNITTEST_CONFIG, BERT2BERTConfig
from official.nlp.nhnet.models import Bert2Bert, get_bert2bert_layers

bert2bert_config_dict = UNITTEST_CONFIG.copy()
bert2bert_config_dict["len_title"] = 32
bert2bert_config = BERT2BERTConfig.from_args(**bert2bert_config_dict)
bert_layer, decoder_layer = get_bert2bert_layers(params=bert2bert_config)
bert2bert = Bert2Bert(bert2bert_config, bert_layer, decoder_layer)

inputs = {
    "input_ids": tf.keras.layers.Input((200,), dtype=tf.int32, name="input_ids"),
    "input_mask": tf.keras.layers.Input((200,), dtype=tf.int32, name="input_mask"),
    "segment_ids": tf.keras.layers.Input((200,), dtype=tf.int32, name="segment_ids")
}

output = bert2bert(inputs, mode='predict')
model = tf.keras.Model(inputs=inputs, outputs=[output])
model.save("bert2bert")

you'll get the following stack:

  [...]
  File "C:\dev\ml\OnnxConversionLab\venv\lib\site-packages\keras\saving\saved_model\save_impl.py", line 633, in call_and_return_conditional_losses
    call_output = layer_call(*args, **kwargs)
  File "C:\dev\ml\OnnxConversionLab\venv\lib\site-packages\official\nlp\nhnet\models.py", line 235, in call
    inputs["target_ids"], bias_type="decoder_self")
KeyError: 'target_ids'

4. Expected behavior

Saving should work using only the input needed at serving time.

5. Additional context

Adding:

"target_ids": tf.keras.layers.Input((32,), dtype=tf.int32, name="target_ids")

to the list of inputs make the saving works.

6. System information

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10.0.19042
TensorFlow installed from (source or binary): binary
TensorFlow version (use command below): 2.6.0
Python version: 3.7.6

saberkun commented 3 years ago

bert2bert needs input_ids, input_mask, segment_ids and target_ids to train. You should save the model with all features provided.

If you care about inference and there is no target_ids, you should not use Keras model.save(). Keras does not support None as inputs. Instead, we directly define a tf.Module including the bert2bert core model and save the tf.function using tf.saved_model.save() API. Usually, the seq2seq model is not friendly to Keras assumptions.

google-ml-butler[bot] commented 3 years ago

Are you satisfied with the resolution of your issue? Yes No

LoicDagnas commented 3 years ago

@saberkun in fact I get the exact same error using

tf.saved_model.save(model, "any/path")

at the end of my example above.

tensorflow / models