keras
keras copied to clipboard
Delegate `finalize_variable_values` in `LossScaleOptimizerV3`
This PR delegates finalize_variable_values()
in LossScaleOptimizerV3
to the wrapped optimizer.
This fixes exponential moving averages when used together with mixed precision.
Prior to this change
import tensorflow as tf
tf.keras.mixed_precision.set_global_policy("mixed_float16")
dataset = tf.data.Dataset.from_tensors(([1.0], [1.0])).repeat(100).batch(10)
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
model.compile(loss="mse", optimizer=tf.keras.optimizers.SGD(use_ema=True))
model.fit(dataset, epochs=2)
would throw
/usr/local/lib/python3.7/dist-packages/keras/optimizers/optimizer_experimental/optimizer.py in _overwrite_model_variables_with_average_value(self, var_list)
677 def _overwrite_model_variables_with_average_value(self, var_list):
678 """Overwrite model variables with its moving average."""
--> 679 if len(var_list) != len(self._model_variables_moving_average):
680 raise ValueError(
681 f"The length of model variables ({len(var_list)}) to "
AttributeError: 'LossScaleOptimizerV3' object has no attribute '_model_variables_moving_average'
See this colab notebook.
/cc @reedwm @fchollet would it be possible to also cherry-pick this fix onto r2.11
so it can be included in the 2.11.0 stable release?
@fchollet, This should be cherrypicked to 2.11 release before final.