Repository: tensorflow/addons

Issue:
Failing to restore RectifiedAdam optimizer from checkpoint
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
- TensorFlow version and how it was installed (source or binary): source, nightly
- TensorFlow-Addons version and how it was installed (source or binary): source, nightly
- Python version: 3.6.7
- Is GPU used? (yes/no): yes

**Describe the bug**

When restoring the RectifiedAdam optimizer from a checkpoint, restoring fails with the following error message:
```ValueError: Tensor conversion requested dtype float32 for Tensor with dtype int32: <tf.Tensor: shape=(), dtype=int32, numpy=0>```

The reason is that RectifiedAdam stores `total_steps` as an integer hyper-parameter:
https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/rectified_adam.py#L134

OptimizerV2 however creates variables by calling add_weight without a dtype in `_create_hypers`. In other words, it assumes that all hyper-parameters are floating point parameters.

Storing the `total_steps` parameter as float fixes the problem for me.


Discussion:

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
@@ -103,7 +103,7 @@ def __init__(
                 beyond".
             sma_threshold. A float value.
                 The threshold for simple mean average.
-            total_steps: An integer. Total number of training steps.
+            total_steps: An integer value. Total number of training steps.
                 Enable warmup by setting a positive value.
             warmup_proportion: A floating point value.
                 The proportion of increasing steps.
@@ -131,7 +131,7 @@ 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", int(total_steps))
+        self._set_hyper("total_steps", float(total_steps))
         self._set_hyper("warmup_proportion", warmup_proportion)
         self._set_hyper("min_lr", min_lr)
         self.epsilon = epsilon or tf.keras.backend.epsilon()
@@ -316,7 +316,7 @@ def get_config(self):
                 "sma_threshold": self._serialize_hyperparameter("sma_threshold"),
                 "epsilon": self.epsilon,
                 "amsgrad": self.amsgrad,
-                "total_steps": self._serialize_hyperparameter("total_steps"),
+                "total_steps": int(self._serialize_hyperparameter("total_steps")),
                 "warmup_proportion": self._serialize_hyperparameter(
                     "warmup_proportion"
                 ),


Test patch:
diff --git a/tensorflow_addons/optimizers/tests/rectified_adam_test.py b/tensorflow_addons/optimizers/tests/rectified_adam_test.py
--- a/tensorflow_addons/optimizers/tests/rectified_adam_test.py
+++ b/tensorflow_addons/optimizers/tests/rectified_adam_test.py
@@ -209,3 +209,26 @@ def test_scheduler_serialization():
         "class_name": "InverseTimeDecay",
         "config": wd_scheduler.get_config(),
     }
+
+
+def test_checkpoint_serialization(tmpdir):
+    optimizer = RectifiedAdam()
+    optimizer2 = RectifiedAdam()
+
+    var_0 = tf.Variable([1.0, 2.0], dtype=tf.dtypes.float32)
+    var_1 = tf.Variable([3.0, 4.0], dtype=tf.dtypes.float32)
+
+    grad_0 = tf.constant([0.1, 0.2], dtype=tf.dtypes.float32)
+    grad_1 = tf.constant([0.03, 0.04], dtype=tf.dtypes.float32)
+
+    grads_and_vars = list(zip([grad_0, grad_1], [var_0, var_1]))
+
+    optimizer.apply_gradients(grads_and_vars)
+
+    checkpoint = tf.train.Checkpoint(optimizer=optimizer)
+    checkpoint2 = tf.train.Checkpoint(optimizer=optimizer2)
+    model_path = str(tmpdir / "rectified_adam_chkpt")
+    checkpoint.write(model_path)
+    checkpoint2.read(model_path)
+
+    optimizer2.apply_gradients(grads_and_vars)