How to load model with custom loss?

Question

How to load model with custom loss that subclass tf.keras.losses.Loss?

I defined ContrastiveLoss by subclassing tf.keras.losses.Loss as follows:

import tensorflow as tf
from tensorflow.keras.losses import Loss

class ContrastiveLoss(Loss):
    def __init__(self, alpha, square=True, **kwargs):
        super(ContrastiveLoss, self).__init__(**kwargs)
        self.alpha = alpha
        self.square = square

    def get_dists(self, x, y, square):
        dists = tf.subtract(x, y)
        dists = tf.reduce_sum(tf.square(dists), axis=-1)

        if not square:
            zero_mask = tf.cast(tf.equal(dists, 0.0), tf.float32)
            dists = dists + zero_mask * 1e-16
            dists = tf.sqrt(dists)
            nonzero_mask = 1.0 - zero_mask
            dists = dists * nonzero_mask

        return dists

    def call(self, y_true, y_pred):
        # y_true & y_pred shape == (N, #embed), for N mini-batch
        # y_true[:, 0] == (N)
        if len(y_true.shape) == 2: y_true= y_true[:, 0]

        positive_mask = tf.cast(tf.equal( tf.expand_dims(y_true, 0), tf.expand_dims(y_true, 1) ), tf.float32)
        negative_mask = tf.subtract(1.0, positive_mask)   

        all_dists = self.get_dists(tf.expand_dims(y_pred, 1), tf.expand_dims(y_pred, 0), self.square)
        positive_loss = tf.multiply( positive_mask, all_dists )
        negative_loss = tf.multiply( negative_mask, tf.maximum(tf.subtract(self.alpha, all_dists), 0.) )
        contrastive_loss = tf.add( positive_loss, negative_loss )

        valid_doublet_mask = tf.cast( tf.greater(contrastive_loss, 1e-16), tf.float32)
        num_valid_doublet = tf.reduce_sum(valid_doublet_mask)
        contrastive_loss  = tf.reduce_sum( contrastive_loss ) / (num_valid_doublet + 1e-16)       

        return contrastive_loss 

    def get_config(self):
        config = super(ContrastiveLoss, self).get_config()
        config.update({'alpha' : self.alpha, 
                       'square' : self.square})
        return config  

I can train and save model with it.

However when I load model as follows, I get error messages.

load_model(model_path, custom_objects={'ContrastiveLoss' : ContrastiveLoss})

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-af42cd2404e1> in <module>()
----> 1 load_model(model_path, custom_objects={'ContrastiveLoss' : ContrastiveLoss})

/tensorflow-2.1.0/python3.6/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(

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/saving/saved_model/load.py in load(path, compile)
     97     if training_config is not None:
     98       model.compile(**saving_utils.compile_args_from_training_config(
---> 99           training_config))
    100   # pylint: disable=protected-access
    101 

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/saving/saving_utils.py in compile_args_from_training_config(training_config, custom_objects)
    232   loss_config = training_config['loss']  # Deserialize loss class.
    233   if isinstance(loss_config, dict) and 'class_name' in loss_config:
--> 234     loss_config = losses.get(loss_config)
    235   loss = nest.map_structure(
    236       lambda obj: custom_objects.get(obj, obj), loss_config)

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/losses.py in get(identifier)
   1184     return deserialize(identifier)
   1185   if isinstance(identifier, dict):
-> 1186     return deserialize(identifier)
   1187   elif callable(identifier):
   1188     return identifier

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/losses.py in deserialize(name, custom_objects)
   1173       module_objects=globals(),
   1174       custom_objects=custom_objects,
-> 1175       printable_module_name='loss function')
   1176 
   1177 

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    290     config = identifier
    291     (cls, cls_config) = class_and_config_for_serialized_keras_object(
--> 292         config, module_objects, custom_objects, printable_module_name)
    293 
    294     if hasattr(cls, 'from_config'):

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    248     cls = module_objects.get(class_name)
    249     if cls is None:
--> 250       raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    251 
    252   cls_config = config['config']

ValueError: Unknown loss function: ContrastiveLoss

It is strange that if I use custom loss "function", there is no error during load_model(.)

But in this case, using "subclass" of Loss, error occurs.

Complete code

Answer 1

If what javad suggests

Have you tried using an object instead of the class name, meaning load_model(model_path, custom_objects={'ContrastiveLoss' : ContrastiveLoss(...)}) where ... has all the parameters for your loss like alpha,...?

Doesn't work and you want to make inference only, then try using:

tf.keras.models.load_model("<model_path>", compile=False)

Hope it helps.