Failing to load saved TFBertModel
Author: sourabhXIII
Pull request: no
Labels: name: wontfix

TF version: 2.2.0-rc1
transformers version: 2.7.0


`import tensorflow as tf`
`import transformers`
`print(tf.__version__)`
`print(transformers.__version__)`
`MAX_LEN = 10`
`model_path = 'saved_model/temp_model'`
`ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)`
`mask = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)`
`token_type_ids = tf.keras.layers.Input((MAX_LEN,), dtype=tf.int32)`
`base_model = transformers.TFBertModel.from_pretrained("bert-base-cased"`
                `, output_hidden_states=False)`
`base_output = base_model([ids, mask, token_type_ids])`
`seq_out, _ = base_output[0], base_output[1]`
`base_model.trainable = False`
`model = tf.keras.models.Model(inputs=[ids, mask, token_type_ids], outputs=[seq_out])`
`model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])`
`print(model.summary())`
`model.save(model_path)`
`model = tf.keras.models.load_model(model_path)`

Model load fails with the following error:


Traceback (most recent call last):
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 378, in assert_same_structure
    expand_composites)
TypeError: The two structures don't have the same nested structure.

First structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}

Second structure: type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')]

More specifically: The two namedtuples don't have the same sequence type. First structure type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')} has type dict, while second structure type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')] has type list

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "temp.py", line 29, in <module>
    model = tf.keras.models.load_model(model_path)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/save.py", line 190, in load_model
    return saved_model_load.load(filepath, compile)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 116, in load
    model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 604, in load_internal
    export_dir)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 188, in __init__
    super(KerasObjectLoader, self).__init__(*args, **kwargs)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/saved_model/load.py", line 123, in __init__
    self._load_all()
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 215, in _load_all
    self._finalize_objects()
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 506, in _finalize_objects
    _finalize_saved_model_layers(layers_revived_from_saved_model)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 677, in _finalize_saved_model_layers
    inputs = infer_inputs_from_restored_call_function(call_fn)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 921, in infer_inputs_from_restored_call_function
    spec = nest.map_structure(common_spec, spec, spec2)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 611, in map_structure
    expand_composites=expand_composites)
  File "/Users/sourabhmaity/anaconda3/lib/python3.7/site-packages/tensorflow/python/util/nest.py", line 385, in assert_same_structure
    % (str(e), str1, str2))
TypeError: The two structures don't have the same nested structure.

First structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')}

Second structure: type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')]

More specifically: The two namedtuples don't have the same sequence type. First structure type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='input_ids')} has type dict, while second structure type=list str=[TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/0'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/1'), TensorSpec(shape=(None, 10), dtype=tf.int32, name='inputs/2')] has type list
Entire first structure:
{'input_ids': .}
Entire second structure:
[., ., .]

Comments:
Facing a similar issue with tf 2.2.0-rc3.
```
def get_model(lr=0.00001):
    inp_bert = tf.keras.layers.Input(shape=(512), dtype="int32")
    bert = TFBertModel.from_pretrained('bert-base-multilingual-cased')(inp_bert)[0]
    doc_encodings = tf.squeeze(bert[:, 0:1, :], axis=1)
    out = tf.keras.layers.Dense(1, activation="sigmoid")(doc_encodings)
    model = tf.keras.Model(inp_bert, out)
    adam = tf.keras.optimizers.Adam(lr=lr)
    model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
    return model
model = get_model()
model.save("model_name",save_format='tf')
model = tf.keras.models.load_model('model_name')
model.summary()
```
Output error is:

```
/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
    383                   "Entire first structure:\n%s\n"
    384                   "Entire second structure:\n%s"
--> 385                   % (str(e), str1, str2))
    386 
    387 

ValueError: The two structures don't have the same nested structure.

First structure: type=TensorSpec str=TensorSpec(shape=(None, 512), dtype=tf.int32, name='inputs')

Second structure: type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}

More specifically: Substructure "type=dict str={'input_ids': TensorSpec(shape=(None, 5), dtype=tf.int32, name='inputs/input_ids')}" is a sequence, while substructure "type=TensorSpec str=TensorSpec(shape=(None, 512), dtype=tf.int32, name='inputs')" is not
Entire first structure:
.
Entire second structure:
{'input_ids': .}
```

change  
`base_output = base_model([ids, mask, token_type_ids])`  
to   
`base_output = base_model.bert([ids, mask, token_type_ids])`
should fix


> 
> 
> change
> `base_output = base_model([ids, mask, token_type_ids])`
> to
> `base_output = base_model.bert([ids, mask, token_type_ids])`
> should fix

Thanks @Souls362 .. solves it.

This worked for me as well with `TFBertModel`, however, I run into the same issue with `TFXLNetModel`. `TFXLNetModel` doesn't seem to have an equivalent to the `.bert` property/attribute. Does anyone know how to solve this when using `TFXLNetModel`?

For `TFXLNetModel` this would be the `.transformer` attribute, as you can see [here](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_tf_xlnet.py#L1127)

@LysandreJik thank you! That works perfectly

@LysandreJik How about `TFOpenAIGPTLMHeadModel` ? I use `.transformer` attribute, but the output shape become `[None, None, 768]`, while the original output shape of `TFOpenAIGPTLMHeadModel` is `[None, None, 13088]`. How to solve it? Thanks a lot!

Well the `transformer` attribute is the transformer in itself, which has a hidden size of 768. The LM head model has an additional head which is the embedding matrix of, which has a size of 13088.

Yes, I think so too. So how can I save the whole model?

@Souls362  you are the greatest! I looked way too long for this. 
@huggingface folks, please add an extra detailed example on serialization in TF2. 
There is seriously some clear documentation missing there.

> change
> `base_output = base_model([ids, mask, token_type_ids])`
> to
> `base_output = base_model.bert([ids, mask, token_type_ids])`
> should fix

one tip for TFBertSequenceClassification:  base_model.bert([ids, mask, token_type_ids])[1]

> > change
> > `base_output = base_model([ids, mask, token_type_ids])`
> > to
> > `base_output = base_model.bert([ids, mask, token_type_ids])`
> > should fix
> 
> one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]

What is the difference of 0 and 1 in the brackets?

> > > change
> > > `base_output = base_model([ids, mask, token_type_ids])`
> > > to
> > > `base_output = base_model.bert([ids, mask, token_type_ids])`
> > > should fix
> > 
> > 
> > one tip for TFBertSequenceClassification: base_model.bert([ids, mask, token_type_ids])[1]
> 
> What is the difference of 0 and 1 in the brackets?

[TFBertModel documentation](https://huggingface.co/transformers/model_doc/bert.html#transformers.TFBertModel)

model returns sequence output and pooled output (for classification)

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


> change
> `base_output = base_model([ids, mask, token_type_ids])`
> to
> `base_output = base_model.bert([ids, mask, token_type_ids])`
> should fix

best answer ever!

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_92500/3480800436.py in <module>
----> 1 model_eval = model.evaluate(
      2     dataset_test,
      3     use_multiprocessing=True,
      4     return_dict=True)

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py in evaluate(self, x, y, batch_size, verbose, sample_weight, steps, callbacks, max_queue_size, workers, use_multiprocessing, return_dict)
   1387             with trace.Trace('test', step_num=step, _r=1):
   1388               callbacks.on_test_batch_begin(step)
-> 1389               tmp_logs = self.test_function(iterator)
   1390               if data_handler.should_sync:
   1391                 context.async_wait()

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    826     tracing_count = self.experimental_get_tracing_count()
    827     with trace.Trace(self._name) as tm:
--> 828       result = self._call(*args, **kwds)
    829       compiler = "xla" if self._experimental_compile else "nonXla"
    830       new_tracing_count = self.experimental_get_tracing_count()

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    869       # This is the first call of __call__, so we have to initialize.
    870       initializers = []
--> 871       self._initialize(args, kwds, add_initializers_to=initializers)
    872     finally:
    873       # At this point we know that the initialization is complete (or less

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    723     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    724     self._concrete_stateful_fn = (
--> 725         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    726             *args, **kwds))
    727 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   2967       args, kwargs = None, None
   2968     with self._lock:
-> 2969       graph_function, _ = self._maybe_define_function(args, kwargs)
   2970     return graph_function
   2971 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3359 
   3360           self._function_cache.missed.add(call_context_key)
-> 3361           graph_function = self._create_graph_function(args, kwargs)
   3362           self._function_cache.primary[cache_key] = graph_function
   3363 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3194     arg_names = base_arg_names + missing_arg_names
   3195     graph_function = ConcreteFunction(
-> 3196         func_graph_module.func_graph_from_py_func(
   3197             self._name,
   3198             self._python_function,

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    988         _, original_func = tf_decorator.unwrap(python_func)
    989 
--> 990       func_outputs = python_func(*func_args, **func_kwargs)
    991 
    992       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    632             xla_context.Exit()
    633         else:
--> 634           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    635         return out
    636 

~/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    975           except Exception as e:  # pylint:disable=broad-except
    976             if hasattr(e, "ag_error_metadata"):
--> 977               raise e.ag_error_metadata.to_exception(e)
    978             else:
    979               raise

TypeError: in user code:

    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1233 test_function  *
        return step_function(self, iterator)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1224 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1217 run_step  **
        outputs = model.test_step(data)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/training.py:1188 test_step
        self.compiled_metrics.update_state(y, y_pred, sample_weight)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/compile_utils.py:387 update_state
        self.build(y_pred, y_true)
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/keras/engine/compile_utils.py:317 build
        self._metrics = nest.map_structure_up_to(y_pred, self._get_metric_objects,
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/util/nest.py:1159 map_structure_up_to
        return map_structure_with_tuple_paths_up_to(
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/util/nest.py:1241 map_structure_with_tuple_paths_up_to
        assert_shallow_structure(
    /home/user/miniconda3/envs/p3/lib/python3.8/site-packages/tensorflow/python/util/nest.py:847 assert_shallow_structure
        raise TypeError(_STRUCTURES_HAVE_MISMATCHING_TYPES.format(

    TypeError: The two structures don't have the same sequence type. Input structure has type <class 'tuple'>, while shallow structure has type <class 'dict'>.
```

```python
def map_tk(X, Y):
  X = tokenizer(
    X,
    max_length=max_len,
    padding='max_length',
    truncation=True,
    return_token_type_ids=False,
    return_tensors="tf")
  
  X = {
    "input_ids": tf.reshape(X["input_ids"], [max_len]),
    "attention_mask": tf.reshape(X["attention_mask"], [max_len])
  }
  Y = {
    "y1": to_categorical(Y["y1"], num_classes=11), 
    "y2": to_categorical(Y["y2"], num_classes=4)
    }
  return X, Y


def gen_data(df: pd.DataFrame):
  def gen():
    for _, row in df.iterrows():
      d = {
        "X": row["content"], 
        "Y": {
          "y1": row["y1"],
          "y2": row["y2"]
        }
      }
      yield map_tk(d["X"], d["Y"])
  
  return gen


output_signature = (
  {
    "input_ids": tf.TensorSpec(shape=(150,), dtype=tf.int32),
    "attention_mask": tf.TensorSpec(shape=(150,), dtype=tf.int32)
  },
  {
    "institution": tf.TensorSpec(shape=(11,), dtype=tf.int32),
    "laws_nature": tf.TensorSpec(shape=(4,), dtype=tf.int32)
  })


def build_dataset(df: pd.DataFrame, shuffle_size=0):
  ds = tf.data.Dataset.from_generator(
    gen_data(df),
    output_signature=output_signature)

  if shuffle_size > 0:
    ds = ds.shuffle(buffer_size=shuffle_size)

  return ds.batch(batch_size=batch_size).prefetch(1)


dataset_train = build_dataset(train, 25600)
dataset_valid = build_dataset(valid)
dataset_test = build_dataset(test)
```

This is the inputs of my model, is there any workaround for `electra`?
```python
input_ids = Input(shape=(max_len,), name="input_ids", dtype="int32")
attention_mask = Input(shape=(max_len,), name="attention_mask", dtype="int32")
inputs = {"input_ids": input_ids, "attention_mask": attention_mask}

X = pretrained(inputs)["hidden_states"][-3:-1]
```

I had the same problem, this link solved the problem for me -> [link](https://stackoverflow.com/questions/73557769/valueerror-unknown-layer-tfbertmodel-please-ensure-this-object-is-passed-to-t)
also I saved model with Pickle and got a problem loading that too, I couldn't solve that if anyone knows how.

Thanks