tensorflow: keras.models.load_model() fails when the model uses a keras.losses.Loss subclass

i find the reason, in my code: loss_object=tf.losses.SparseCategoricalCrossentropy() model.complie(loss=loss_object, optimizer=“sgd”)

it raise the error.

the i change my code to model.complie(loss=“sparse_categorical_crossentropy”, optimizer=“sgd”)

it is ok

Same issue on the stable 2.0.0 release 😦

This is fixed TF 2.0 nightly ‘2.0.0-dev20190718’ Thanks!

Hi there, I am still seeing the same issue.

Below is my model code.

class HardNegativeContrastiveLoss(keras.losses.Loss):
    """
    Args:
      margin: Control the similarity threshold between two vectors
      name: Name of the loss function.
    """
    def __init__(self, margin=1.0, name='hard_negative_contrastive_loss'):
        super(HardNegativeContrastiveLoss, self).__init__(name=name)
        self.margin = margin

    def call(self, y_true, y_pred):
        """y_true are all positive, they are passed in but not used,
        y_pred contains the embeddings of two sources, the two sources have dimensions
        (batch_size, embedding_dim), they have been concatenated in a new axis to form
        the y_pred: (batch_size, embdding_dim, 2)
        """
        
        embedding1 = y_pred[:, :, 0]
        embedding2 = y_pred[:, :, 1]
        embedding1_norm = K.l2_normalize(embedding1, axis=1)
        embedding1_norm_t = K.transpose(embedding1_norm)

        embedding2_norm = K.l2_normalize(embedding2, axis=1)

        dot_product = K.dot(embedding2_norm, embedding1_norm_t)
        distance = 2.0 - 2.0 * dot_product

        distance = tf.maximum(distance, 0.0)
        mask = tf.cast(tf.equal(distance, 0.0), 'float32') # correction to avoid sqrt(0.0) that would cause gradient of sqrt to be infinite
        distance = distance + mask * 1e-16
        distance = tf.sqrt(distance)
        distance = distance * (1.0 - mask)  # correct distance for epsilon added

        positive_score = tf.linalg.diag_part(distance)

        num_sample = K.shape(embedding1)[0]
        indices_equal = tf.cast(tf.eye(num_sample), tf.bool)
        indices_not_equal = tf.logical_not(indices_equal)
        mask_indices_not_equal = tf.cast(indices_not_equal, 'float32')

        negative_scores = tf.multiply(mask_indices_not_equal, distance)

        hardest_negative_score_embedding1 = tf.reduce_min(negative_scores + 4.0 * tf.cast(indices_equal, 'float32'), axis=1)
        hardest_negative_score_embedding2 = tf.reduce_min(negative_scores + 4.0 * tf.cast(indices_equal, 'float32'), axis=0)

        score = K.mean(2.0 * positive_score + tf.maximum(self.margin - hardest_negative_score_embedding1, 0)  + tf.maximum(self.margin - hardest_negative_score_embedding2, 0))
        return score

def basic_embedding_model_cosine(query_vocab_size, title_vocab_size, embedding_dim = 128, use_mask_zero=True):
    """model that evaluates similarity of query and title embeddings using dot product
    uses word tokenization"""

    print('query vocab size:', query_vocab_size)
    print('title vocab size:', title_vocab_size)
    mirrored_strategy = tf.distribute.MirroredStrategy()
    with mirrored_strategy.scope():
        query_input = keras.Input(name="query_input", shape=(constants.MAX_QUERY_WORD_LENGTH,))
        query_embedding = layers.Embedding(input_dim=query_vocab_size + 1,
                                    output_dim=embedding_dim,
                                    mask_zero=use_mask_zero,
                                    name="query_embedding")(query_input)
        query_embedding_normalized = layers.BatchNormalization(name='query_embedding_normalized')(query_embedding)
        query_embedding_vector = layers.GlobalAveragePooling1D(name='query_mean')(query_embedding_normalized)

        title_input = keras.Input(shape=(constants.MAX_TITLE_WORD_LENGTH,), name='title_input')
        title_embedding = layers.Embedding(input_dim=title_vocab_size + 1,
                                   output_dim=embedding_dim,
                                   mask_zero=use_mask_zero,
                                   name="title_embedding")(title_input)
        title_embedding_normalized = layers.BatchNormalization(name='title_embedding_normalized')(title_embedding)
        title_embedding_vector = layers.GlobalAveragePooling1D(name='title_mean')(title_embedding_normalized)

        loss_output = layers.concatenate([tf.expand_dims(query_embedding_vector,-1), tf.expand_dims(title_embedding_vector,-1)],name='loss_output')
        similarity_output = layers.Dot(axes=1,normalize=True,name='similarity_output')([title_embedding_vector, query_embedding_vector])

        model = keras.Model(inputs=[query_input, title_input], outputs=[loss_output, similarity_output])

        optimizer = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
        custom_loss = HardNegativeContrastiveLoss(margin=1.0)

        model.compile(
            optimizer=optimizer,
            loss={'loss_output': custome_loss},
            metrics={'similarity_output':binary_accuracy}
        )

        return model

After successfully train and save the model, I come across with ‘Unknown entries in loss dictionary’ error. Full error message as follows:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-55-8b153b1e1e9a> in <module>()
----> 1 unreplicated_model = tf.keras.models.load_model(model_path)
      2 # unreplicated_model.compile(
      3 #     optimizer=Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0),
      4 #     loss={'output_1':HardNegativeContrastiveLoss(margin=1.0)}
      5 # )

~/anaconda3/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/save.py in load_model(filepath, custom_objects, compile)
    148   if isinstance(filepath, six.string_types):
    149     loader_impl.parse_saved_model(filepath)
--> 150     return saved_model_load.load(filepath, compile)
    151 
    152   raise IOError(

~/anaconda3/lib/python3.6/site-packages/tensorflow_core/python/keras/saving/saved_model/load.py in load(path, compile)
     91     if model._training_config is not None:  # pylint: disable=protected-access
     92       model.compile(**saving_utils.compile_args_from_training_config(
---> 93           model._training_config))  # pylint: disable=protected-access
     94 
     95   return model

~/anaconda3/lib/python3.6/site-packages/tensorflow_core/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    455     self._self_setattr_tracking = False  # pylint: disable=protected-access
    456     try:
--> 457       result = method(self, *args, **kwargs)
    458     finally:
    459       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

~/anaconda3/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py in compile(self, optimizer, loss, metrics, loss_weights, sample_weight_mode, weighted_metrics, target_tensors, distribute, **kwargs)
    334     # Prepare list of loss functions, same size of model outputs.
    335     self.loss_functions = training_utils.prepare_loss_functions(
--> 336         self.loss, self.output_names)
    337 
    338     target_tensors = self._process_target_tensor_for_compile(target_tensors)

~/anaconda3/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_utils.py in prepare_loss_functions(loss, output_names)
   1339   """
   1340   if isinstance(loss, collections_abc.Mapping):
-> 1341     generic_utils.check_for_unexpected_keys('loss', loss, output_names)
   1342     loss_functions = []
   1343     for name in output_names:

~/anaconda3/lib/python3.6/site-packages/tensorflow_core/python/keras/utils/generic_utils.py in check_for_unexpected_keys(name, input_dict, expected_values)
    589     raise ValueError('Unknown entries in {} dictionary: {}. Only expected '
    590                      'following keys: {}'.format(name, list(unknown),
--> 591                                                  expected_values))
    592 
    593 

ValueError: Unknown entries in loss dictionary: ['loss_output']. Only expected following keys: ['output_1', 'output_2']

Anyone can help?

I am using tf.version=‘2.0.0’ on Ubuntu 16.6.

Can confirm changing the loss function to a string removes this error.