Model save issues with custom optimizer - DPKerasSGDOptimizer

[avssridhar]

Problem
A keras sequential Model was compiled using DPKerasSGD optimizer and saved. But when the saved model is loaded it throws the below error

"Unknown optimizer: 'DPOptimizerClass'. Please ensure you are using a keras.utils.custom_object_scope and that this object is included in the scope. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details"

Expected Behaviour
Since the model was already compiled using the optimizer it should run perfectly fine when fitted with the data.

Whereas the same works when the optimizer is replaced. Please provide a solution on how to save the model with this custom optimizer.

Standalone code to reproduce the issue

#building the model
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid')])

#compiling the model
model.compile(
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True,reduction=tf.losses.Reduction.NONE),
optimizer=tensorflow_privacy.DPKerasSGDOptimizer(l2_norm_clip=1,noise_multiplier=2,learning_rate=0.01),
metrics=[
tf.keras.metrics.BinaryAccuracy(name='accuracy'),
tf.keras.metrics.Precision(name='precision'),
tf.keras.metrics.Recall(name='recall')
])

#saving the model
model.save('my_model')

#loading the model
keras.models.load_model('mymodel')

#throws the error

[sushreebarsa]

@avssridhar

FYI, There are two formats you can use to save an entire model to disk: the TensorFlow SavedModel format, and the older Keras H5 format. The recommended format is SavedModel. It is the default when you use model.save().

You can switch to the H5 format by:

Passing save_format='h5' to save().
Passing a filename that ends in .h5 or .keras to save().
For more information, you may refer to this guide as well.
Alternative approach can be followed as below;

my_loaded_model = tf.keras.models.load_model('my_models_name.h5', custom_objects={'KerasLayer':hub.KerasLayer , 'AdamWeightDecay': optimizer})

In order to expedite the trouble-shooting process, please provide the complete code snippet to reproduce the issue reported here. While executing your code, I faced a different outcome as shown in the gist.
Thank you!

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please reopen if you'd like to work on this further.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[Ruairi-Grant]

Hi, how did you resolve this? I have a similar issue where my optimizer is DPKerasSGDOptimizer and I am training like this:

checkpoint = ModelCheckpoint(
    "best_model.h5", monitor="val_loss", mode="min", verbose=1, save_best_only=True
)

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=EPOCHS,
    callbacks=[early_stopping, checkpoint],
)

best_model = load_model("best_model.h5")

I get the same error "Unknown optimizer: 'DPOptimizerClass'. Please ensure you are using a keras.utils.custom_object_scope and that this object is included in the scope. See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object for details" when loading best_model.h5