Unable to restore custom object of type _tf_keras_metric currently using the HammingLoss metric from TensorFlow Addons module

Viewed 2149

I have created and trained a TensorFlow model using the HammingLoss metric from TensorFlow addons. Thus, it's not a custom metric that I have created on my own. I use a callbacks function with the methords ModelCheckpoint() and EarlyStopping to save the best weights of the best model and stop model training at a given threshold repsectively. When I save the model checkpoint I serialize the whole model structure (similar to model.save()), istead of model.save_weights(), which would have saved only the model weights (more about ModelCheckpoint here).

TL;DR: Here is a colab notebook with the code I post below in case you want to skip this.

The model I have trained is saved in GoogleDrive in the link here. To load the specific model I use the following code:

neural_network_parameters = {}

#======================================================================
#           PARAMETERS THAT DEFINE THE NEURAL NETWORK STRUCTURE       =
#======================================================================

neural_network_parameters['model_loss'] = tf.keras.losses.BinaryCrossentropy(from_logits=False, name='binary_crossentropy')

neural_network_parameters['model_metric'] = [tfa.metrics.HammingLoss(mode="multilabel", name="hamming_loss"),
                                             tfa.metrics.F1Score(17, average="micro", name="f1_score_micro"),
                                             tfa.metrics.F1Score(17, average=None, name="f1_score_none"),
                                             tfa.metrics.F1Score(17, average="macro", name="f1_score_macro"),
                                             tfa.metrics.F1Score(17, average="weighted", name="f1_score_weighted")]

"""Initialize the hyper parameters tuning the model using Tensorflow's hyperparameters module"""
HP_HIDDEN_UNITS = hp.HParam('batch_size', hp.Discrete([32]))
HP_EMBEDDING_DIM = hp.HParam('embedding_dim', hp.Discrete([50]))
HP_LEARNING_RATE = hp.HParam('learning_rate', hp.Discrete([0.001])) # Adam default: 0.001, SGD default: 0.01, RMSprop default: 0.001....0.1 to be removed
HP_DECAY_STEPS_MULTIPLIER = hp.HParam('decay_steps_multiplier', hp.Discrete([10]))

METRIC_ACCURACY = "hamming_loss"

dependencies = {
    'hamming_loss': tfa.metrics.HammingLoss(mode="multilabel", name="hamming_loss"),
    'attention': attention(return_sequences=True)
}

def import_trained_keras_model(model_index, method, decay_steps_mode, optimizer_name, hparams):
    """Load the model"""
    
    training_date="2021-02-27"
    model_path_structure=f"{folder_path_model_saved}/{initialize_notebbok_variables.saved_model_name}_{hparams[HP_EMBEDDING_DIM]}dim_{hparams[HP_HIDDEN_UNITS]}batchsize_{hparams[HP_LEARNING_RATE]}lr_{hparams[HP_DECAY_STEPS_MULTIPLIER]}decaymultiplier_{training_date}"
    model_imported=load_model(f"{model_path_structure}", custom_objects=dependencies)

    if optimizer_name=="adam":
        optimizer = optimizer_adam_v2(hparams)
    
    elif optimizer_name=="sgd":
        optimizer = optimizer_sgd_v1(hparams, "step decay")

    else:
        optimizer = optimizer_rmsprop_v1(hparams)

    model_imported.compile(optimizer=optimizer,
                            loss=neural_network_parameters['model_loss'],
                            metrics=neural_network_parameters['model_metric'])

    print(f"Model {model_index} is loaded successfully\n")
    
    return model_imported

Calling the function import trained keras model

"""Now that the functions have been created it's time to import each trained classifier from the selected dictionary of hyper parameters, calculate the evaluation metric per model and finally serialize the scores dataframe for later use."""
list_models=[] #a list to store imported models
model_optimizer="adam"

for batch_size in HP_HIDDEN_UNITS.domain.values:
  for embedding_dim in HP_EMBEDDING_DIM.domain.values:
      for learning_rate in HP_LEARNING_RATE.domain.values:
          for decay_steps_multiplier in HP_DECAY_STEPS_MULTIPLIER.domain.values:
              hparams = {
                  HP_HIDDEN_UNITS: batch_size,
                  HP_EMBEDDING_DIM: embedding_dim,
                  HP_LEARNING_RATE: learning_rate,
                  HP_DECAY_STEPS_MULTIPLIER: decay_steps_multiplier
                }
              print(f"\n{len(list_models)+1}/{(len(HP_HIDDEN_UNITS.domain.values)*len(HP_EMBEDDING_DIM.domain.values)*len(HP_LEARNING_RATE.domain.values)*len(HP_DECAY_STEPS_MULTIPLIER.domain.values))}")
              print({h.name: hparams[h] for h in hparams},'\n')
              
              model_object=import_trained_keras_model(len(list_models)+1, "import custom trained model", "on", model_optimizer, hparams)
              list_models.append(model_object)

When I call the function I get the following error

ValueError: Unable to restore custom object of type _tf_keras_metric currently. Please make sure that the layer implements get_configand from_config when saving. In addition, please use the custom_objects arg when calling load_model().

It's strange that I get this error since the model metric to compile the NN is from a built in method of TensorFlow and NOT some sort of a custom metric that I developed myself.

I have searched also this thread in GitHub which closed without explaining the root of the problem.

[UPDATE]--Found a temporary solution

I managed to successfully import the model by turning the compile argument to False in order to re-compile the model imported inside the function.

So I did smth like model_imported=load_model(f"{model_path_structure}", custom_objects=dependencies, compile=False).

This action produced the following result:

WARNING:tensorflow:Unable to restore custom metric. Please ensure that the layer implements get_config and from_config when saving. In addition, please use the custom_objects arg when calling load_model().

Model 1 is loaded successfully.

So TensorFlow still cannot understand that HammingLoss is not a custom metric but rather a metric imported from Tensorflow Addons. However, despite the warning the model loaded successfully.

0 Answers
Related