加载保存的 BertClassifer model

Question

I am using the following example from this colab notebook: https://github.com/tensorflow/models/blob/master/official/colab/fine_tuning_bert.ipynb

It saves a fine-tuned model using the model.save() functionality.

I am trying to load that same model using tf.keras.models.load_model() but am getting the following error:

KeyError: 'name'
KeyError                                  Traceback (most recent call last)
in engine
----> 1 model = tf.keras.models.load_model('./saved_model/test')

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/save.py in load_model(filepath, custom_objects, compile)
    183     if isinstance(filepath, six.string_types):
    184       loader_impl.parse_saved_model(filepath)
--> 185       return saved_model_load.load(filepath, compile)
    186 
    187   raise IOError(

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in load(path, compile)
    114   # TODO(kathywu): Add saving/loading of optimizer, compiled losses and metrics.
    115   # TODO(kathywu): Add code to load from objects that contain all endpoints
--> 116   model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
    117 
    118   # pylint: disable=protected-access

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py in load_internal(export_dir, tags, loader_cls)
    605       loader = loader_cls(object_graph_proto,
    606                           saved_model_proto,
--> 607                           export_dir)
    608       root = loader.get(0)
    609       if isinstance(loader, Loader):

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in __init__(self, *args, **kwargs)
    186     self._models_to_reconstruct = []
    187 
--> 188     super(KerasObjectLoader, self).__init__(*args, **kwargs)
    189 
    190     # Now that the node object has been fully loaded, and the checkpoint has

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py in __init__(self, object_graph_proto, saved_model_proto, export_dir)
    121       self._concrete_functions[name] = _WrapperFunction(concrete_function)
    122 
--> 123     self._load_all()
    124     self._restore_checkpoint()
    125 

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _load_all(self)
    207     # loaded from config may create variables / other objects during
    208     # initialization. These are recorded in `_nodes_recreated_from_config`.
--> 209     self._layer_nodes = self._load_layers()
    210 
    211     # Load all other nodes and functions.

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _load_layers(self)
    307         continue
    308 
--> 309       layers[node_id] = self._load_layer(proto.user_object, node_id)
    310 
    311     for node_id, proto in metric_list:

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _load_layer(self, proto, node_id)
    333     # Detect whether this object can be revived from the config. If not, then
    334     # revive from the SavedModel instead.
--> 335     obj, setter = self._revive_from_config(proto.identifier, metadata, node_id)
    336     if obj is None:
    337       obj, setter = revive_custom_object(proto.identifier, metadata)

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _revive_from_config(self, identifier, metadata, node_id)
    350     else:
    351       obj = (
--> 352           self._revive_graph_network(metadata, node_id) or
    353           self._revive_layer_from_config(metadata, node_id))
    354 

/home/cdsw/.local/lib/python3.6/site-packages/tensorflow/python/keras/saving/saved_model/load.py in _revive_graph_network(self, metadata, node_id)
    386     else:
    387       model = models_lib.Functional(
--> 388           inputs=[], outputs=[], name=config['name'])
    389 
    390     # Record this model and its layers. This will later be used to reconstruct

KeyError: 'name'

Can someone please advise what I am doing wrong or if I should be doing something else to load this model? Thanks!

Answer 1

As it stated, name is not defined. model = tf.keras.models.load_model('./saved_model/test') leads to the PATH (folder) ./saved_model/test , not the PATH of the model.

Answer 2

I've documented my experience of that KeyError (and the solution I found) here:

https://github.com/tensorflow/models/issues/8692#issuecomment-727033061

It was due to a bad implementation of a custom layer's get_config() that didn't use the parent's base config, that does contain the name key, necessary for keras reload.

Answer 3

Since the provided notebook is no longer available I will assume that you were using this (oficial?) BERT implementation for thensorflow since I was having the same issue with it.

---

TL;DR: You can't load the model due to missing proper get_config methods in the subclassed model and (possibly) this layer . You need to edit the BERT encoder class (if this is a possibility for you).

---

I came across this while trying to modify the BERT model and apply the Masked Language Modeling training. I found that I was not able to load the checkpoints due to the mentioned KeyError: 'name' . I found three issues:

1. After a lot of testing I found that several layers where producing a warning during saving, the Warning stated Custom mask layers require a config and must override... I fixed the layers in my implementation, however the problem persisted.
2. Later I realized that the KeyError: 'name' came from the model's get_config implementation which was not yielding the tf.keras.Model original config along with the new config dict . In this particular case it was not due to missing layer configuration (as it is often the case). I patched the subclassed model get_config method in my implementation:

def get_config(self):
    # Implement get_config to enable serialization. This is NOT optional.
    base_config = super(LayoutEncoder, self).get_config()
    return dict(list(self._config.items()) + list(base_config.items()))

3. Solving the KeyError: 'name' (using the code in point 2.) created a new and more horrible error (also during loading/saving). More than 500 lines of errors/warnings containing WARNING:tensorflow:Unresolved object in checkpoint... and Inconsistent references when loading the checkpoint into this object graph... appeared and the loading still failed. I found that this was caused by the way that the BERT model was keeping track of the layers (lines 484-491). These layers were also serialized during save. I'm not sure why these members are added to the encoder class, however by removing them I was able to save and load the model without any other problem.

Sorry for the long history but, imho, this was the only way to explain it...