Save / Load Tensorflow Keras model for Prediction only

user_2532768 0
Viewed 1764

I have a tensorflow keras model with custom losses. After trianing, I want to store the model using model.save(path) and in another python script load the model for prediction only using model = tf.keras.models.load_model(model_path, compile=False).

Without compile=False, tensorflow will complain about the missing loss function. Calling model.prediction however will result in a Model object has no attribute 'loss' error message. I would like to call model.predict without needing to specify the loss again.

Is there a solution to save/load a tf.keras.Model without the custom loss to using the model for prediction?

Code

Since it was asked, the model trains on multiple outputs / losses and I define the losses with lambdas to capture weightings etc. This looks like this: 

losses = [lambda y_true, y_pred: util.weighted_mse_loss(y_true, y_pred, tf.square(gain_weight)), 
        lambda y_true, y_pred: util.weighted_mse_loss(y_true, y_pred, tf.square(Rd_weight)), 
        lambda y_true, y_pred: util.pole_zero_loss(y_true, y_pred, r_weight, w_weight),
        lambda y_true, y_pred: util.pole_zero_loss(y_true, y_pred, r_weight, w_weight)]


model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=10E-4),
    loss=losses)
tensorflow2.0 tf.keras tensorflow
1 Answers

As mentioned here keras has tf.keras.models.load_model(filepath, custom_objects=None, compile=True) function. If you have custom loss, you can use:

model = keras.models.load_model('model.h5', custom_objects={'my_loss': my_loss})

Where 'my_loss' is your loss' name, and my_loss is the function.

user_115891830
Related