Repository: keras-team/keras

Issue:
Chrome timeline is broken for TensorFlow backend
Chrome timeline is very useful to profile a Keras model, we can get the execution time for each node in the TF graph. Keras has already supported it since https://github.com/keras-team/keras/pull/6693 . But it seems this feature is broken since Keras 2.2. Run the following code to reproduce this bug:
```
import keras
from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Dropout
from keras.optimizers import RMSprop

import tensorflow as tf
from tensorflow.python.client import timeline

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train.reshape(60000, 784).astype('float32')
y_train = keras.utils.to_categorical(y_train, 10)

model = Sequential()
model.add(Dense(512, activation='relu', input_shape=(784,)))
model.add(Dropout(0.2))
model.add(Dense(512, activation='relu'))
model.add(Dropout(0.2))
model.add(Dense(10, activation='softmax'))

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()

model.compile(loss='categorical_crossentropy',
              optimizer=RMSprop(),
              metrics=['accuracy'],
              options=run_options,
              run_metadata=run_metadata)

history = model.fit(x_train, y_train,
                    batch_size=128,
                    epochs=1)

trace = timeline.Timeline(step_stats=run_metadata.step_stats)
with open('/tmp/timeline.json', 'w') as f:
    f.write(trace.generate_chrome_trace_format())
```
This is the exception:
```
Using TensorFlow backend.
Traceback (most recent call last):
  File "test.py", line 72, in <module>
    epochs=1)
  File "build/bdist.macosx-10.13-intel/egg/keras/engine/training.py", line 1016, in fit
  File "build/bdist.macosx-10.13-intel/egg/keras/engine/training.py", line 516, in _make_train_function
  File "build/bdist.macosx-10.13-intel/egg/keras/backend/tensorflow_backend.py", line 2705, in function
  File "build/bdist.macosx-10.13-intel/egg/keras/backend/tensorflow_backend.py", line 2552, in __init__
ValueError: ('Some keys in session_kwargs are not supported at this time: %s', ['run_metadata', 'options'])
```
It seems ```run_metadata``` and ```options``` arguments are not supported by ```K.Function```.


Discussion:

Resolving patch:
diff --git a/keras/backend/tensorflow_backend.py b/keras/backend/tensorflow_backend.py
--- a/keras/backend/tensorflow_backend.py
+++ b/keras/backend/tensorflow_backend.py
@@ -2545,7 +2545,10 @@ def __init__(self, inputs, outputs,
         # (since the outputs of fetches are never returned).
         # This requires us to wrap fetches in `identity` ops.
         self.fetches = [tf.identity(x) for x in self.fetches]
-        self.session_kwargs = session_kwargs
+        # self.session_kwargs is used for _legacy_call
+        self.session_kwargs = session_kwargs.copy()
+        self.run_options = session_kwargs.pop('options', None)
+        self.run_metadata = session_kwargs.pop('run_metadata', None)
         if session_kwargs:
             raise ValueError('Some keys in session_kwargs are not '
                              'supported at this '
@@ -2593,6 +2596,9 @@ def _make_callable(self, feed_arrays, feed_symbols, symbol_vals, session):
             callable_opts.fetch.append(x.name)
         # Handle updates.
         callable_opts.target.append(self.updates_op.name)
+        # Handle run_options.
+        if self.run_options:
+            callable_opts.run_options.CopyFrom(self.run_options)
         # Create callable.
         callable_fn = session._make_callable_from_options(callable_opts)
         # Cache parameters corresponding to the generated callable, so that
@@ -2643,7 +2649,10 @@ def _call(self, inputs):
                                 feed_symbols,
                                 symbol_vals,
                                 session)
-        fetched = self._callable_fn(*array_vals)
+        if self.run_metadata:
+            fetched = self._callable_fn(*array_vals, run_metadata=self.run_metadata)
+        else:
+            fetched = self._callable_fn(*array_vals)
         return fetched[:len(self.outputs)]
 
     def _legacy_call(self, inputs):
@@ -2673,6 +2682,16 @@ def __call__(self, inputs):
                         'supported with sparse inputs.')
                 return self._legacy_call(inputs)
 
+            # callable generated by Session._make_callable_from_options accepts
+            # `run_metadata` keyword argument since TF 1.10
+            if (self.run_metadata and
+                    StrictVersion(tf.__version__.split('-')[0]) < StrictVersion('1.10.0')):
+                if py_any(is_tensor(x) for x in inputs):
+                    raise ValueError(
+                        'In order to feed symbolic tensors to a Keras model and set '
+                        '`run_metadata`, you need tensorflow 1.10 or higher.')
+                return self._legacy_call(inputs)
+
             return self._call(inputs)
         else:
             if py_any(is_tensor(x) for x in inputs):


Test patch:
diff --git a/tests/keras/backend/backend_test.py b/tests/keras/backend/backend_test.py
--- a/tests/keras/backend/backend_test.py
+++ b/tests/keras/backend/backend_test.py
@@ -528,6 +528,31 @@ def test_function_tf_feed_dict(self):
         assert output == [21.]
         assert K.get_session().run(fetches=[x, y]) == [30., 40.]
 
+    @pytest.mark.skipif(K.backend() != 'tensorflow',
+                        reason='Uses the `options` and `run_metadata` arguments.')
+    def test_function_tf_run_options_with_run_metadata(self):
+        from tensorflow.core.protobuf import config_pb2
+        x_placeholder = K.placeholder(shape=())
+        y_placeholder = K.placeholder(shape=())
+
+        run_options = config_pb2.RunOptions(output_partition_graphs=True)
+        run_metadata = config_pb2.RunMetadata()
+        # enable run_options.
+        f = K.function(inputs=[x_placeholder, y_placeholder],
+                       outputs=[x_placeholder + y_placeholder],
+                       options=run_options,
+                       run_metadata=run_metadata)
+        output = f([10., 20.])
+        assert output == [30.]
+        assert len(run_metadata.partition_graphs) > 0
+        # disable run_options.
+        f = K.function(inputs=[x_placeholder, y_placeholder],
+                       outputs=[x_placeholder + y_placeholder],
+                       run_metadata=run_metadata)
+        output = f([10., 20.])
+        assert output == [30.]
+        assert len(run_metadata.partition_graphs) == 0
+
     @pytest.mark.skipif(K.backend() != 'tensorflow',
                         reason='Uses the `string` type for a tensor.')
     def test_function_tf_string_input(self):