Repository: tensorflow/addons

Issue:
MovingAverage cannot be used with callbacks.ModelCheckpoint
**Describe the feature and the current behavior/state.**
Mostly, I save the model at the point of highest val_accuracy by `callbacks.ModelCheckpoint`, but the `MovingAverage` optimizer doesn't seem to support `callbacks.ModelCheckpoint`, i.e., we cannot save the model with moving average in the middle of training processing!

I think we should add MovingAverage into callbacks, and save the moving average model in the middle of training.

**Relevant information**
- Are you willing to contribute it (yes/no):
- Are you willing to maintain it going forward? (yes/no):
- Is there a relevant academic paper? (if so, where):
- Is there already an implementation in another framework? (if so, where): https://github.com/alno/kaggle-allstate-claims-severity/blob/master/keras_util.py
- Was it part of tf.contrib? (if so, where):

**Which API type would this fall under (layer, metric, optimizer, etc.)**
callbacks

**Who will benefit with this feature?**
all tf.keras users

**Any other info.**



Discussion:
cc @Squadrick 
@wuliytTaotao Could you post the code that caused the error? I've been trying it out on `tfa-nightly`, and I can't find any error. 
@Squadrick Here is an example, using the model.fit() and callbacks.ModelCheckpoint to save the weights at the point of best validation accuracy, but moving average cannot be added into this processing. I think it's best to implement this function as a callback function.
```
import tensorflow as tf
import os
from tensorflow_addons.optimizers.moving_average import MovingAverage

os.environ["CUDA_VISIBLE_DEVICES"] = '1'

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, enable=True)
print(gpus)

cifar10 = tf.keras.datasets.cifar10

(x_train, y_train), (x_test, y_test) = cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(16, (3, 3), padding='same', activation='relu', input_shape=(32, 32, 3)),
    tf.keras.layers.AveragePooling2D(),
    tf.keras.layers.Conv2D(32, (3, 3), padding='same', activation='relu'),
    tf.keras.layers.AveragePooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])


optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
optimizer = MovingAverage(optimizer)

model.compile(optimizer=optimizer,
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5, validation_split=0.1,
          # weights.h5 cannot be saved with moving average,
          # for "average_op" function of MovingAverage cannot be applied during training
          callbacks=[tf.keras.callbacks.ModelCheckpoint('weights.h5', monitor="val_accuracy",
                                                        save_weights_only=True, save_best_only=True)])

model.load_weights("weights.h5")
model.evaluate(x_test, y_test, verbose=2)
```

Resolving patch:
diff --git a/tensorflow_addons/callbacks/__init__.py b/tensorflow_addons/callbacks/__init__.py
--- a/tensorflow_addons/callbacks/__init__.py
+++ b/tensorflow_addons/callbacks/__init__.py
@@ -14,5 +14,6 @@
 # ==============================================================================
 """Additional callbacks that conform to Keras API."""
 
+from tensorflow_addons.callbacks.average_model_checkpoint import AverageModelCheckpoint
 from tensorflow_addons.callbacks.time_stopping import TimeStopping
 from tensorflow_addons.callbacks.tqdm_progress_bar import TQDMProgressBar
diff --git a/tensorflow_addons/callbacks/average_model_checkpoint.py b/tensorflow_addons/callbacks/average_model_checkpoint.py
new file mode 100644
--- /dev/null
+++ b/tensorflow_addons/callbacks/average_model_checkpoint.py
@@ -0,0 +1,83 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+import tensorflow as tf
+from typeguard import typechecked
+from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper
+
+
+class AverageModelCheckpoint(tf.keras.callbacks.ModelCheckpoint):
+    r"""The callback that should be used with optimizers that extend
+    AverageWrapper, i.e., MovingAverage and StochasticAverage optimizers.
+    It saves and, optionally, assigns the averaged weights.
+
+    Args:
+        update_weights: If True, assign the moving average weights
+            to the model, and save them. If False, keep the old
+            non-averaged weights, but the saved model uses the
+            average weights.
+
+        See `tf.keras.callbacks.ModelCheckpoint` for the other args.
+    """
+
+    @typechecked
+    def __init__(
+        self,
+        update_weights: bool,
+        filepath: str,
+        monitor: str = "val_loss",
+        verbose: int = 0,
+        save_best_only: bool = False,
+        save_weights_only: bool = False,
+        mode: str = "auto",
+        save_freq: str = "epoch",
+        **kwargs
+    ):
+        self.update_weights = update_weights
+        super().__init__(
+            filepath,
+            monitor,
+            verbose,
+            save_best_only,
+            save_weights_only,
+            mode,
+            save_freq,
+            **kwargs,
+        )
+
+    def set_model(self, model):
+        if not isinstance(model.optimizer, AveragedOptimizerWrapper):
+            raise TypeError(
+                "AverageModelCheckpoint is only used when training"
+                "with MovingAverage or StochasticAverage"
+            )
+        return super().set_model(model)
+
+    def _save_model(self, epoch, logs):
+        assert isinstance(self.model.optimizer, AveragedOptimizerWrapper)
+
+        if self.update_weights:
+            self.model.optimizer.assign_average_vars(self.model.variables)
+            return super()._save_model(epoch, logs)
+        else:
+            # Note: `model.get_weights()` gives us the weights (non-ref)
+            # whereas `model.variables` returns references to the variables.
+            non_avg_weights = self.model.get_weights()
+            self.model.optimizer.assign_average_vars(self.model.variables)
+            # result is currently None, since `super._save_model` doesn't
+            # return anything, but this may change in the future.
+            result = super()._save_model(epoch, logs)
+            self.model.set_weights(non_avg_weights)
+            return result


Test patch: