tfa.optimizers.MultiOptimizer exception when mixed_precision is enabled.

[umarov90]

System information

• OS Platform and Distribution: Linux Ubuntu 16.04
• TensorFlow version and how it was installed (source or binary): 2.8.2 pip
• TensorFlow-Addons version and how it was installed (source or binary): 0.17.0 pip
• Python version: 3.7
• Is GPU used? (yes/no): yes

Describe the bug

tfa.optimizers.MultiOptimizer exception when mixed_precision is enabled. The following code runs without a problem when mixed_precision is not enabled. Using mixed precision I get the following exception: TypeError: apply_gradients() takes 2 positional arguments but 3 were given

Please see the full log bellow. Could you tell me how I could fix the issue? Thank you.

Code to reproduce the issue

import numpy as np import tensorflow as tf import tensorflow_addons as tfa from tensorflow.keras.layers import Dense, Input, Flatten from tensorflow.keras.models import Model from tensorflow.keras import mixed_precision mixed_precision.set_global_policy('mixed_float16')

inputs = Input(shape=(10, 4)) x = inputs x = Dense(128, activation="relu", name="layer1")(x) x = Dense(64, activation="relu", name="layer2")(x) x = Flatten()(x) outputs = Dense(1, activation="relu", dtype=tf.float32)(x) our_model = Model(inputs, outputs, name="our_model")

optimizers_and_layers = [(tf.keras.optimizers.Adam(learning_rate=0.001), our_model.get_layer("layer1")), (tf.keras.optimizers.Adam(learning_rate=0.005), our_model.get_layer("layer2"))] optimizer = tfa.optimizers.MultiOptimizer(optimizers_and_layers) our_model.compile(optimizer=optimizer, loss="mse") X = np.random.rand(50, 10, 4) Y = np.random.rand(50, 1) our_model.fit(X, Y)

Other info / logs

Traceback (most recent call last): File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler raise e.with_traceback(filtered_tb) from None File "xxxxxxxx/venv/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 1147, in autograph_handler raise e.ag_error_metadata.to_exception(e) TypeError: in user code:

File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/engine/training.py", line 1021, in train_function  *
    return step_function(self, iterator)
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/engine/training.py", line 1010, in step_function  **
    outputs = model.distribute_strategy.run(run_step, args=(data,))
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/engine/training.py", line 1000, in run_step  **
    outputs = model.train_step(data)
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/engine/training.py", line 863, in train_step
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py", line 532, in minimize
    return self.apply_gradients(grads_and_vars, name=name)
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/mixed_precision/loss_scale_optimizer.py", line 640, in apply_gradients
    do_not_apply_fn)
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/mixed_precision/loss_scale_optimizer.py", line 637, in apply_fn
    return self._apply_gradients(grads, wrapped_vars, name)
File "xxxxxxxx/venv/lib/python3.7/site-packages/keras/mixed_precision/loss_scale_optimizer.py", line 673, in _apply_gradients
    experimental_aggregate_gradients=False)

TypeError: apply_gradients() takes 2 positional arguments but 3 were given

python-BaseException

[m274d]

System information

• OS Platform and Distribution: Ubuntu 20.04.3 LTS
• TensorFlow version and how it was installed (source or binary): 2.8.0
• TensorFlow-Addons version and how it was installed (source or binary): 0.16.1
• Python version: 3.8
• Is GPU used? (yes/no): yes

I got exactly the same issue when using tfa.optimizers.MultiOptimizer with mixed_precision in a self-defined training loop. The only difference it I further wrap the MultiOptimizer inside the LossScaleOptimizer.

optimizer = mixed_precision.LossScaleOptimizer(optimizer)

The training loop follow exactly the same as the official mixed_precision training guide:

def train_step(x, y):
  with tf.GradientTape() as tape:
    predictions = model(x)
    loss = loss_object(y, predictions)
    scaled_loss = optimizer.get_scaled_loss(loss)
  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)
  gradients = optimizer.get_unscaled_gradients(scaled_gradients)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  return loss

[isarandi]

This was fixed in this commit in keras: https://github.com/keras-team/keras/commit/013adbae4d35448c49ccf36938cdf9a684a8ddd8

[bhack]

Can we close this?

[m274d]

This was fixed in this commit in keras: keras-team/keras@013adba

Thank you! I will upgrade my keras version to try again. I think we can close it for now.