Use tf.Keras with Tensorflow optimizer

Question

I'm trying to use tf.keras with the new AdamW optimizer in tensorflow and am running into issues. A toy version of the code is as follows:

from tensorflow.contrib.opt import AdamWOptimizer
from tensorflow.python.keras.optimizers import TFOptimizer

model = Sequential()
model.add(Dense(2, activation="tanh", input_shape=(3,)))

tfopt = AdamWOptimizer(weight_decay=0.1, learning_rate=.004)
optimizer = TFOptimizer(tfopt)

model.compile(optimizer=optimizer, loss='mean_squared_error')
model.fit(np.random.random((5, 3)),
          np.random.random((5, 2)),
          epochs=5, batch_size=5)

Error is as follows:

../python3.6/site-packages/tensorflow/python/keras/engine/training.py:1605: in fit
    validation_steps=validation_steps)
../python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py:153: in fit_loop
    outs = f(ins)
../python3.6/site-packages/tensorflow/python/keras/backend.py:2978: in __call__
    run_metadata=self.run_metadata)
../python3.6/site-packages/tensorflow/python/client/session.py:1399: in __call__
    run_metadata_ptr)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <tensorflow.python.framework.errors_impl.raise_exception_on_not_ok_status object at 0x11ecde550>
type_arg = None, value_arg = None, traceback_arg = None

    def __exit__(self, type_arg, value_arg, traceback_arg):
      try:
        if c_api.TF_GetCode(self.status.status) != 0:
          raise _make_specific_exception(
              None, None,
              compat.as_text(c_api.TF_Message(self.status.status)),
>             c_api.TF_GetCode(self.status.status))
E             tensorflow.python.framework.errors_impl.FailedPreconditionError: Attempting to use uninitialized value training/TFOptimizer/beta2_power
E                [[{{node training/TFOptimizer/beta2_power/read}} = Identity[T=DT_FLOAT, _class=["loc:@training/TFOptimizer/AdamW/Assign"], _device="/job:localhost/replica:0/task:0/device:CPU:0"](training/TFOptimizer/beta2_power)]]

../python3.6/site-packages/tensorflow/python/framework/errors_impl.py:526: FailedPreconditionError

Answer 1

Turns out TFOptimizer won't work in python 3.6 but it does work in 2.7.

However, you don't actually need to use TFOptimizer. Plugging the tensorflow optimizer AdamWOptimizer directly into the optimizer argument of fit runs perfectly.