TensorFlow 2.0 tf.keras API Eager模式与图形模式

Question

In TensorFlow <2 the training function for a DDPG actor could be concisely implemented using tf.keras.backend.function as follows:

critic_output = self.critic([self.actor(state_input), state_input])
actor_updates = self.optimizer_actor.get_updates(params=self.actor.trainable_weights,
                                                 loss=-tf.keras.backend.mean(critic_output))
self.actor_train_on_batch = tf.keras.backend.function(inputs=[state_input], 
                                                      outputs=[self.actor(state_input)],
                                                      updates=actor_updates)

Then during each training step calling self.actor_train_on_batch([np.array(state_batch)]) would compute the gradients and perform the updates.

However running that on TF 2.0 gives the following error due to eager mode being on by default:

actor_updates = self.optimizer_actor.get_updates(params=self.actor.trainable_weights, loss=-tf.keras.backend.mean(critic_output))
File "C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\keras\optimizer_v2\optimizer_v2.py", line 448, in get_updates
    grads = self.get_gradients(loss, params)
File "C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\keras\optimizer_v2\optimizer_v2.py", line 361, in get_gradients
    grads = gradients.gradients(loss, params)
File "C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\ops\gradients_impl.py", line 158, in gradients
    unconnected_gradients)
File "C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\ops\gradients_util.py", line 547, in _GradientsHelper
    raise RuntimeError("tf.gradients is not supported when eager execution "
RuntimeError: tf.gradients is not supported when eager execution is enabled. Use tf.GradientTape instead.

As expected, disabling eager execution via tf.compat.v1.disable_eager_execution() fixes the issue.

However I don't want to disable eager execution for everything - I would like to use purely the 2.0 API. The exception suggests using tf.GradientTape instead of tf.gradients but that's an internal call.

Question: What is the appropriate way of computing -tf.keras.backend.mean(critic_output) in graph mode (in TensorFlow 2.0)?

Answer 1

As far as I understood, your critic_output is just a TensorFlow tensor, so you can just use tf.math.reduce_mean operation. And it'll work in a TensorFlow session, not in imperative style. Ie this will return an operation to be evaluated in a TensorFlow session.

import tensorflow as tf
import numpy as np

inp = tf.placeholder(dtype=tf.float32)
mean_op = tf.math.reduce_mean(inp)

with tf.Session() as sess:
    print(sess.run(mean_op, feed_dict={inp: np.ones(10)}))
    print(sess.run(mean_op, feed_dict={inp: np.random.randn(10)}))

It'll evaluate in something like:

1.0
-0.002577734

Answer 2

So, first of all you error is related to the fact that optimizer.get_updates() is designed for graph mode as it does include the K.gradients() needed to get the gradients tensors and then apply the Keras optimizer-based update to the trainable variables of the model using the K.function . Secondly, in terms of eager-mode-or-not soundness the cost function loss=-tf.keras.backend.mean(critic_output) has no flows. What you should to is get rid of your graph mode code and stick to the native 2.0 eager mode. Based on your code the training should look like:

def train_method(self, state_input):
  with tf.GradientTape() as tape:
    critic_output = self.critic([self.actor(state_input), state_input])
    loss=-tf.keras.backend.mean(critic_output)
  grads = tape.gradient(loss, params=self.actor.trainable_variables)
  # now please note that self.optimizer_actor must have apply_gradients 
  # so it should be tf.train.OptimizerName...
  self.optimizer_actor.apply_gradients(zip(grads, self.actor.trainable_variables))