Repository: tensorflow/addons

Issue:
Lookahead bug when loading saved model
**System information**
- OS Platform and Distribution: Windows 10
- TensorFlow version and how it was installed (source or binary): TF 2.1 from pip 
- TensorFlow-Addons version and how it was installed (source or binary): 0.8.3
- Python version: python 3.7
- Is GPU used? (yes/no): yes cuda 10.1

**Describe the bug**

My model uses TFA: Lookahead with Radam (ranger)
Trained it: ok
Saved it: ok
Impossible to load it: 

```
Traceback (most recent call last):
  File "Boucle_IA.py", line 24, in <module>
    model = tf.keras.models.load_model('inference_complete-2020-03-18-Mobilenet.h5', custom_objects={'Lookahead': ranger})
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\saving\save.py", line 146, in load_model
    return hdf5_format.load_model_from_hdf5(filepath, custom_objects, compile)
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\saving\hdf5_format.py", line 184, in load_model_from_hdf5
    training_config, custom_objects))
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\saving\saving_utils.py", line 229, in compile_args_from_training_config
    optimizer_config, custom_objects=custom_objects)
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\optimizers.py", line 819, in deserialize
    printable_module_name='optimizer')
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\utils\generic_utils.py", line 303, in deserialize_keras_object
    list(custom_objects.items())))
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_addons\optimizers\lookahead.py", line 185, in from_config
    custom_objects=custom_objects,
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\optimizers.py", line 819, in deserialize
    printable_module_name='optimizer')
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\utils\generic_utils.py", line 303, in deserialize_keras_object
    list(custom_objects.items())))
  File "C:\Program Files\Python37\lib\site-packages\tensorflow_core\python\keras\optimizer_v2\optimizer_v2.py", line 733, in from_config
    return cls(**config)
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 809, in wrapper
    check_argument_types(memo)
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 670, in check_argument_types
    raise exc from None
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 668, in check_argument_types
    check_type(description, value, expected_type, memo)
  File "C:\Program Files\Python37\lib\site-packages\typeguard\__init__.py", line 598, in check_type
    format(argname, qualified_name(expected_type), qualified_name(value)))
TypeError: type of argument "total_steps" must be int; got float instead
```



**Code to reproduce the issue**

on definition : 
```python
radam = tfa.optimizers.RectifiedAdam(
    lr=1e-3,
    total_steps=10000,
    warmup_proportion=0.1,
    min_lr=1e-5,
)
ranger = tfa.optimizers.Lookahead(radam, sync_period=6, slow_step_size=0.5)
```

on load : 
```python
import numpy as np 
import tensorflow as tf 
import tensorflow_addons 
import os
import cv2 

from pathlib import Path

# Constantes

print('loading model...')
radam = tfa.optimizers.RectifiedAdam(lr=1e-3, total_steps=10000, warmup_proportion=0.1, min_lr=1e-5)

ranger = tfa.optimizers.Lookahead(radam, sync_period=6, slow_step_size=0.5, name='Lookahead',)

model = tf.keras.models.load_model('inference_complete-2020-03-18-Mobilenet.h5', custom_objects={'Lookahead': ranger})
print('model loaded')
```

Tips: When I tried to open H5 copy of the saved model: in total_steps, value = 10000.0


Discussion:
cc @cyberzhg

Resolving patch:
diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py
--- a/tensorflow_addons/optimizers/rectified_adam.py
+++ b/tensorflow_addons/optimizers/rectified_adam.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Rectified Adam (RAdam) optimizer."""
+import warnings
 
 import tensorflow as tf
 from tensorflow_addons.utils.types import FloatTensorLike
@@ -79,7 +80,10 @@ def __init__(
         weight_decay: FloatTensorLike = 0.0,
         amsgrad: bool = False,
         sma_threshold: FloatTensorLike = 5.0,
-        total_steps: int = 0,
+        # float for total_steps is here to be able to load models created before
+        # https://github.com/tensorflow/addons/pull/1375 was merged. It should be
+        # removed for Addons 0.11.
+        total_steps: Union[int, float] = 0,
         warmup_proportion: FloatTensorLike = 0.1,
         min_lr: FloatTensorLike = 0.0,
         name: str = "RectifiedAdam",
@@ -123,7 +127,16 @@ def __init__(
         self._set_hyper("decay", self._initial_decay)
         self._set_hyper("weight_decay", weight_decay)
         self._set_hyper("sma_threshold", sma_threshold)
-        self._set_hyper("total_steps", float(total_steps))
+        if isinstance(total_steps, float):
+            warnings.warn(
+                "The parameter `total_steps` passed to the __init__ of RectifiedAdam "
+                "is a float. This behavior is deprecated and in Addons 0.11, this "
+                "will raise an error. Use an int instead. If you get this message "
+                "when loading a model, save it again and the `total_steps` parameter "
+                "will automatically be converted to a int.",
+                DeprecationWarning,
+            )
+        self._set_hyper("total_steps", int(total_steps))
         self._set_hyper("warmup_proportion", warmup_proportion)
         self._set_hyper("min_lr", min_lr)
         self.epsilon = epsilon or tf.keras.backend.epsilon()


Test patch:
diff --git a/tensorflow_addons/optimizers/rectified_adam_test.py b/tensorflow_addons/optimizers/rectified_adam_test.py
--- a/tensorflow_addons/optimizers/rectified_adam_test.py
+++ b/tensorflow_addons/optimizers/rectified_adam_test.py
@@ -172,5 +172,14 @@ def test_get_config(self):
         self.assertEqual(config["total_steps"], 0)
 
 
+def test_serialization():
+    optimizer = RectifiedAdam(
+        lr=1e-3, total_steps=10000, warmup_proportion=0.1, min_lr=1e-5,
+    )
+    config = tf.keras.optimizers.serialize(optimizer)
+    new_optimizer = tf.keras.optimizers.deserialize(config)
+    assert new_optimizer.get_config() == optimizer.get_config()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))