keras icon indicating copy to clipboard operation
keras copied to clipboard

Delegate `finalize_variable_values` in `LossScaleOptimizerV3`

Open lgeiger opened this issue 2 years ago • 1 comments

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?

lgeiger avatar Nov 08 '22 12:11 lgeiger

@fchollet, This should be cherrypicked to 2.11 release before final.

qlzh727 avatar Nov 08 '22 17:11 qlzh727