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:
bot-of-gabrieldemarmiesse: @cyberzhg 

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.

olesalscheider: The tests fail because `total_steps` is int in the constructor. I can fix that, but I am really not sure if this is the right way to handle this, or if this is unintended behavior from tensorflow.
bhack: It is a little bit confusing as a float.
olesalscheider: > It is a little bit confusing as a float.

That's what I thought. But if the constructor takes an int but stores it as a float internally, the tests fail because the optimizer can't be restored.
bhack: Mhh.. I've not run this but can we use a `_get_hyper` with the `dtype` before the set_hyper?
google-cla[bot]: All (the pull request submitter and all commit authors) CLAs are signed, **but** one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project.  Please have them confirm that by leaving a comment that contains only `@googlebot I consent.` in this pull request.

*Note to project maintainer:* There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent.  In those cases, you can manually confirm consent of the commit author(s), and set the `cla` label to `yes` (if enabled on your project).

ℹ️ **Googlers: [Go here](https://goto.google.com/prinfo/https%3A%2F%2Fgithub.com%2Ftensorflow%2Faddons%2Fpull%2F2369) for more info**.

<!-- need_author_consent -->
google-cla[bot]: All (the pull request submitter and all commit authors) CLAs are signed, **but** one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project.  Please have them confirm that by leaving a comment that contains only `@googlebot I consent.` in this pull request.

*Note to project maintainer:* There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent.  In those cases, you can manually confirm consent of the commit author(s), and set the `cla` label to `yes` (if enabled on your project).

ℹ️ **Googlers: [Go here](https://goto.google.com/prinfo/https%3A%2F%2Fgithub.com%2Ftensorflow%2Faddons%2Fpull%2F2369) for more info**.

<!-- need_author_consent -->
bhack: @googlebot I consent

Patch:
diff --git a/.github/workflows/backport.yml b/.github/workflows/backport.yml
index 7c74784..59dad26 100644
--- a/.github/workflows/backport.yml
+++ b/.github/workflows/backport.yml
@@ -12,7 +12,7 @@ jobs:
     steps:
       - name: Backport Bot
         if: github.event.pull_request.merged && ( ( github.event.action == 'closed' && contains( join( github.event.pull_request.labels.*.name ), 'backport') ) || contains( github.event.label.name, 'backport' ) )
-        uses: Gaurav0/backport@d69fd1d # Version 1.0.24
+        uses: Gaurav0/backport@d69fd1d2469762a7b4007f671857e4f94deed0af # Version 1.0.24
         with:
           bot_username: bot-of-gabrieldemarmiesse
           bot_token: 1353d990cdb8b8ceb1b73d301dce83cc0da3db29
diff --git a/.github/workflows/release-drafter.yml b/.github/workflows/release-drafter.yml
index f6b4ff0..9158d7a 100644
--- a/.github/workflows/release-drafter.yml
+++ b/.github/workflows/release-drafter.yml
@@ -10,7 +10,7 @@ jobs:
   update_release_draft:
     runs-on: ubuntu-latest
     steps:
-      - uses: release-drafter/release-drafter@74e7c42 # Version 5.7.0
+      - uses: release-drafter/release-drafter@74e7c423dafbb406c9c18b1638334f67a7c891c3 # Version 5.7.0
         with:
           config-name: release-template.yml
         env:
diff --git a/README.md b/README.md
index 8a1a157..362b137 100644
--- a/README.md
+++ b/README.md
@@ -81,7 +81,7 @@ what it was tested against.
 | TensorFlow Addons | TensorFlow | Python  |
 |:----------------------- |:---|:---------- |
 | tfa-nightly | 2.3, 2.4 | 3.6, 3.7, 3.8 | 
-| tensorflow-addons-0.12.0 | 2.3, 2.4 |3.6, 3.7, 3.8 |
+| tensorflow-addons-0.12.1 | 2.3, 2.4 |3.6, 3.7, 3.8 |
 | tensorflow-addons-0.11.2 | 2.2, 2.3 |3.5, 3.6, 3.7, 3.8 |
 | tensorflow-addons-0.10.0 | 2.2 |3.5, 3.6, 3.7, 3.8 |
 | tensorflow-addons-0.9.1 | 2.1, 2.2 |3.5, 3.6, 3.7 |
@@ -108,7 +108,7 @@ is compiled differently. A typical example of this would be `conda`-installed Te
 | TensorFlow Addons | TensorFlow | Compiler  | cuDNN | CUDA | 
 |:----------------------- |:---- |:---------|:---------|:---------|
 | tfa-nightly | 2.4 | GCC 7.3.1 | 8.0 | 11.0 |
-| tensorflow-addons-0.12.0 | 2.4  | GCC 7.3.1 | 8.0 | 11.0 |
+| tensorflow-addons-0.12.1 | 2.4  | GCC 7.3.1 | 8.0 | 11.0 |
 | tensorflow-addons-0.11.2 | 2.3  | GCC 7.3.1 | 7.6 | 10.1 |
 | tensorflow-addons-0.10.0 | 2.2  | GCC 7.3.1 | 7.6 | 10.1 |
 | tensorflow-addons-0.9.1 | 2.1  | GCC 7.3.1 | 7.6 | 10.1 |
diff --git a/tensorflow_addons/optimizers/cyclical_learning_rate.py b/tensorflow_addons/optimizers/cyclical_learning_rate.py
index 7da41d0..9d82c24 100644
--- a/tensorflow_addons/optimizers/cyclical_learning_rate.py
+++ b/tensorflow_addons/optimizers/cyclical_learning_rate.py
@@ -66,7 +66,7 @@ class CyclicalLearningRate(tf.keras.optimizers.schedules.LearningRateSchedule):
             maximal_learning_rate: A scalar `float32` or `float64` `Tensor` or
                 a Python number.  The maximum learning rate.
             step_size: A scalar `float32` or `float64` `Tensor` or a
-                Python number. Step size.
+                Python number. Step size denotes the number of training iterations it takes to get to maximal_learning_rate.
             scale_fn: A function. Scheduling function applied in cycle
             scale_mode: ['cycle', 'iterations']. Mode to apply during cyclic
                 schedule
@@ -154,7 +154,7 @@ class TriangularCyclicalLearningRate(CyclicalLearningRate):
             maximal_learning_rate: A scalar `float32` or `float64` `Tensor` or
                 a Python number.  The maximum learning rate.
             step_size: A scalar `float32` or `float64` `Tensor` or a
-                Python number. Step size.
+                Python number. Step size denotes the number of training iterations it takes to get to maximal_learning_rate
             scale_mode: ['cycle', 'iterations']. Mode to apply during cyclic
                 schedule
             name: (Optional) Name for the operation.
@@ -223,7 +223,7 @@ class Triangular2CyclicalLearningRate(CyclicalLearningRate):
             maximal_learning_rate: A scalar `float32` or `float64` `Tensor` or
                 a Python number.  The maximum learning rate.
             step_size: A scalar `float32` or `float64` `Tensor` or a
-                Python number. Step size.
+                Python number. Step size denotes the number of training iterations it takes to get to maximal_learning_rate
             scale_mode: ['cycle', 'iterations']. Mode to apply during cyclic
                 schedule
             name: (Optional) Name for the operation.
@@ -294,7 +294,7 @@ class ExponentialCyclicalLearningRate(CyclicalLearningRate):
             maximal_learning_rate: A scalar `float32` or `float64` `Tensor` or
                 a Python number.  The maximum learning rate.
             step_size: A scalar `float32` or `float64` `Tensor` or a
-                Python number. Step size.
+                Python number. Step size denotes the number of training iterations it takes to get to maximal_learning_rate
             scale_mode: ['cycle', 'iterations']. Mode to apply during cyclic
                 schedule
             gamma: A scalar `float32` or `float64` `Tensor` or a
diff --git a/tensorflow_addons/optimizers/rectified_adam.py b/tensorflow_addons/optimizers/rectified_adam.py
index 881840a..c24c555 100644
--- a/tensorflow_addons/optimizers/rectified_adam.py
+++ b/tensorflow_addons/optimizers/rectified_adam.py
@@ -103,7 +103,7 @@ class RectifiedAdam(tf.keras.optimizers.Optimizer):
                 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 @@ class RectifiedAdam(tf.keras.optimizers.Optimizer):
         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 @@ class RectifiedAdam(tf.keras.optimizers.Optimizer):
                 "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
index 040e5e6..ae003dd 100644
--- 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)