Add is_legacy_optimizer to optimizer config to keep saving/loading consistent. (#16856)

PiperOrigin-RevId: 463928027

Co-authored-by: Chen Qian <[email protected]>

Author: qlzh727

keras/mixed_precision/loss_scale_optimizer.py

@@ -929,7 +929,9 @@ def from_config(cls, config, custom_objects=None):
                 )
             config["inner_optimizer"] = config.pop("optimizer")
         inner_optimizer = optimizers.deserialize(
-            config["inner_optimizer"], custom_objects=custom_objects
+            config["inner_optimizer"],
+            custom_objects=custom_objects,
+            use_legacy_optimizer=True,
         )
         del config["inner_optimizer"]
         return cls(inner_optimizer, **config)
@@ -1366,7 +1368,9 @@ def get_config(self):
     def from_config(cls, config, custom_objects=None):
         config = config.copy()  # Make a copy, since we mutate config
         inner_optimizer = optimizers.deserialize(
-            config["inner_optimizer"], custom_objects=custom_objects
+            config["inner_optimizer"],
+            custom_objects=custom_objects,
+            use_legacy_optimizer=False,
         )
         del config["inner_optimizer"]
         return cls(inner_optimizer, **config)

keras/optimizers/__init__.py

@@ -124,6 +124,12 @@ def deserialize(config, custom_objects=None, **kwargs):
     )
 
     use_legacy_optimizer = kwargs.pop("use_legacy_optimizer", True)
+    if len(config["config"]) > 0:
+        # If the optimizer config is not empty, then we use the value of
+        # `is_legacy_optimizer` to override `use_legacy_optimizer`. If
+        # `is_legacy_optimizer` does not exist in config, it means we are
+        # using the legacy optimzier.
+        use_legacy_optimizer = config["config"].get("is_legacy_optimizer", True)
     if (
         tf.__internal__.tf2.enabled()
         and tf.executing_eagerly()

keras/optimizers/optimizer_experimental/optimizer.py

@@ -97,6 +97,7 @@ def _create_iteration_variable(self):
             )
 
     def _process_kwargs(self, kwargs):
+        kwargs.pop("is_legacy_optimizer", None)
         legacy_kwargs = {
             "lr",
             "decay",
@@ -619,6 +620,7 @@ def get_config(self):
             "ema_momentum": self.ema_momentum,
             "ema_overwrite_frequency": self.ema_overwrite_frequency,
             "jit_compile": self.jit_compile,
+            "is_legacy_optimizer": False,
         }
         return config
 

keras/optimizers/optimizer_experimental/optimizer_test.py

@@ -297,6 +297,7 @@ def testGetAndFromConfig(self):
             "use_ema": True,
             "ema_momentum": 0.5,
             "ema_overwrite_frequency": 50,
+            "is_legacy_optimizer": False,
         }
         self.assertDictContainsSubset(expected_config, config)
         restored_optimizer = adam_new.Adam.from_config(config)