[SOLVED] How to use Lazy Adam optimizer in tensorflow 2.0.0

How to use Lazy Adam optimizer in tensorflow 2.0.0

This code doesnt work: it has problem with tf.contrib

model.compile(optimizer=TFOptimizer(tf.contrib.opt.LazyAdamOptimizer()), loss='categorical_crossentropy')

I have tried something with tensorflow_addons.optimizers.LazyAdam() but that does not work either.

Any ideas how to run LazyAdam in tensorflow 2.0.0 ?

PS: only Adam works well as following:

model.compile(optimizer=tf.keras.optimizers.Adam(), loss='categorical_crossentropy')

Solution

import tensorflow_addons as tfa
optimizer = tfa.optimizers.LazyAdam()

tensorflow_addons is an extra functionality for TensorFlow 2.x, but now Tensorflow 2.x is still not very stable, if you are facing with module 'tensorflow_core.keras.utils' has no attribute 'register_keras_serializable', try to update you tensorflow to the latest stable version.