<p dir="auto"><strong>System information</strong></p>
<ul dir="auto">
<li>OS Platform: Win 10</li>
<li>TensorFlow installed from (source or binary): binary</li>
<li>TensorFlow version (use command below): b'v1.13.1-0-g6612da8951' 1.13.1</li>
<li>Python version: 3.6.7</li>
<li>CUDA/cuDNN version: CUDA V10.0.130, cuDNN 7.5</li>
<li>GPU model and memory: RTX 2070 8GB</li>
</ul>
<p dir="auto"><strong>Describe the current behavior</strong><br>
I've convert my keras network to tensorflow eager and this network dosn't learn anything. The network predicts only the same values. The performance is 2-3 times slower as the keras model.<br>
With tensorflow graph without eager execution everything works fine.</p>
<p dir="auto"><strong>Describe the expected behavior</strong><br>
I want to predict actions for various input images in reinforcement learning</p>
<p dir="auto"><strong>Code to reproduce the issue</strong><br>
<strong>Keras model</strong></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
import tensorflow as tf


class KerasTest:
    def __init__(self, state_space, action_space, lr):
        self.state_space = state_space
        self.action_space = action_space
        self.lr = lr

        inputs = Input(shape=(84, 84, 4))
        rewards = Input(shape=(1,))

        x = Conv2D(32, kernel_size=[8, 8], padding='valid', strides=[4, 4], activation=None)(inputs)
        x = BatchNormalization(trainable=True, epsilon=1e-5)(x)
        x = Activation('relu')(x)
        x = Conv2D(64, kernel_size=[4, 4], padding='valid', strides=[2, 2], activation=None)(x)
        x = BatchNormalization(trainable=True, epsilon=1e-5)(x)
        x = Activation('relu')(x)
        x = Conv2D(64, kernel_size=(4, 4), padding='valid', strides=(2, 2), activation=None)(x)
        x = BatchNormalization(trainable=True, epsilon=1e-5)(x)
        x = Activation('relu')(x)
        x = Flatten()(x)
        x = Dense(512, activation='relu')(x)
        x = Dropout(0.5)(x)
        x = Dense(256, activation='relu')(x)
        x = Dropout(0.3)(x)
        x = Dense(128, activation='relu')(x)
        x = Dropout(0.05)(x)
        logits = Dense(self.action_space, activation=None)(x)

        self.model = Model(inputs=[inputs, rewards], outputs=logits)

        def policy_loss(r):
            def loss(labels, logits):
                policy = tf.nn.softmax(logits)
                entropy = tf.nn.softmax_cross_entropy_with_logits_v2(labels=policy, logits=logits)
                log = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels, logits=logits)
                p_loss = log * tf.stop_gradient(r)
                p_loss = p_loss - 0.01 * entropy
                total_loss = tf.reduce_mean(p_loss)
                return total_loss

            return loss

        self.model.compile(optimizer=Adam(lr=lr), loss=policy_loss(rewards))
        self.model.summary()

    def get_probs(self, s):
        s = s[np.newaxis, :]
        probs = self.model.predict([s, np.array([1])])
        probs = probs.squeeze()
        probs = self.softmax(probs)
        return probs

    def softmax(self, x):
        e_x = np.exp(x - np.max(x))
        return e_x / e_x.sum(axis=0)

    def update_policy(self, s, r, a):
        self.model.train_on_batch([s, r], a)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">keras</span>.<span class="pl-s1">layers</span> <span class="pl-k">import</span> <span class="pl-c1">*</span>
<span class="pl-k">from</span> <span class="pl-s1">keras</span>.<span class="pl-s1">models</span> <span class="pl-k">import</span> <span class="pl-v">Model</span>
<span class="pl-k">from</span> <span class="pl-s1">keras</span>.<span class="pl-s1">optimizers</span> <span class="pl-k">import</span> <span class="pl-v">Adam</span>
<span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>


<span class="pl-k">class</span> <span class="pl-v">KerasTest</span>:
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">state_space</span>, <span class="pl-s1">action_space</span>, <span class="pl-s1">lr</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">state_space</span> <span class="pl-c1">=</span> <span class="pl-s1">state_space</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">action_space</span> <span class="pl-c1">=</span> <span class="pl-s1">action_space</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">lr</span> <span class="pl-c1">=</span> <span class="pl-s1">lr</span>

        <span class="pl-s1">inputs</span> <span class="pl-c1">=</span> <span class="pl-v">Input</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">84</span>, <span class="pl-c1">84</span>, <span class="pl-c1">4</span>))
        <span class="pl-s1">rewards</span> <span class="pl-c1">=</span> <span class="pl-v">Input</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,))

        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Conv2D</span>(<span class="pl-c1">32</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>[<span class="pl-c1">8</span>, <span class="pl-c1">8</span>], <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'valid'</span>, <span class="pl-s1">strides</span><span class="pl-c1">=</span>[<span class="pl-c1">4</span>, <span class="pl-c1">4</span>], <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">inputs</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Activation</span>(<span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Conv2D</span>(<span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>[<span class="pl-c1">4</span>, <span class="pl-c1">4</span>], <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'valid'</span>, <span class="pl-s1">strides</span><span class="pl-c1">=</span>[<span class="pl-c1">2</span>, <span class="pl-c1">2</span>], <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Activation</span>(<span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Conv2D</span>(<span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">4</span>, <span class="pl-c1">4</span>), <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'valid'</span>, <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">2</span>, <span class="pl-c1">2</span>), <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Activation</span>(<span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Flatten</span>()(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-c1">512</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-c1">0.5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-c1">256</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-c1">0.3</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-c1">128</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-c1">0.05</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">logits</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-s1">self</span>.<span class="pl-s1">action_space</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">x</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-v">Model</span>(<span class="pl-s1">inputs</span><span class="pl-c1">=</span>[<span class="pl-s1">inputs</span>, <span class="pl-s1">rewards</span>], <span class="pl-s1">outputs</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)

        <span class="pl-k">def</span> <span class="pl-en">policy_loss</span>(<span class="pl-s1">r</span>):
            <span class="pl-k">def</span> <span class="pl-en">loss</span>(<span class="pl-s1">labels</span>, <span class="pl-s1">logits</span>):
                <span class="pl-s1">policy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax</span>(<span class="pl-s1">logits</span>)
                <span class="pl-s1">entropy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax_cross_entropy_with_logits_v2</span>(<span class="pl-s1">labels</span><span class="pl-c1">=</span><span class="pl-s1">policy</span>, <span class="pl-s1">logits</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)
                <span class="pl-s1">log</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax_cross_entropy_with_logits_v2</span>(<span class="pl-s1">labels</span><span class="pl-c1">=</span><span class="pl-s1">labels</span>, <span class="pl-s1">logits</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)
                <span class="pl-s1">p_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">log</span> <span class="pl-c1">*</span> <span class="pl-s1">tf</span>.<span class="pl-en">stop_gradient</span>(<span class="pl-s1">r</span>)
                <span class="pl-s1">p_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">p_loss</span> <span class="pl-c1">-</span> <span class="pl-c1">0.01</span> <span class="pl-c1">*</span> <span class="pl-s1">entropy</span>
                <span class="pl-s1">total_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">reduce_mean</span>(<span class="pl-s1">p_loss</span>)
                <span class="pl-k">return</span> <span class="pl-s1">total_loss</span>

            <span class="pl-k">return</span> <span class="pl-s1">loss</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">compile</span>(<span class="pl-s1">optimizer</span><span class="pl-c1">=</span><span class="pl-v">Adam</span>(<span class="pl-s1">lr</span><span class="pl-c1">=</span><span class="pl-s1">lr</span>), <span class="pl-s1">loss</span><span class="pl-c1">=</span><span class="pl-en">policy_loss</span>(<span class="pl-s1">rewards</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">summary</span>()

    <span class="pl-k">def</span> <span class="pl-en">get_probs</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>):
        <span class="pl-s1">s</span> <span class="pl-c1">=</span> <span class="pl-s1">s</span>[<span class="pl-s1">np</span>.<span class="pl-s1">newaxis</span>, :]
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">predict</span>([<span class="pl-s1">s</span>, <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>])])
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">probs</span>.<span class="pl-en">squeeze</span>()
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">softmax</span>(<span class="pl-s1">probs</span>)
        <span class="pl-k">return</span> <span class="pl-s1">probs</span>

    <span class="pl-k">def</span> <span class="pl-en">softmax</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-s1">e_x</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">exp</span>(<span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">np</span>.<span class="pl-en">max</span>(<span class="pl-s1">x</span>))
        <span class="pl-k">return</span> <span class="pl-s1">e_x</span> <span class="pl-c1">/</span> <span class="pl-s1">e_x</span>.<span class="pl-en">sum</span>(<span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)

    <span class="pl-k">def</span> <span class="pl-en">update_policy</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>, <span class="pl-s1">r</span>, <span class="pl-s1">a</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">train_on_batch</span>([<span class="pl-s1">s</span>, <span class="pl-s1">r</span>], <span class="pl-s1">a</span>)</pre></div>
<p dir="auto"><strong>Eager</strong></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow as tf
from tensorflow.python.keras.layers import *
import numpy as np

tf.enable_eager_execution()
print(tf.executing_eagerly())


class EagerTest:
    def __init__(self, state_space, action_space, lr):
        self.action_space = action_space
        self.lr = lr

        inputs = Input(shape=(84, 84, 4))
        rewards = Input(shape=(1,))

        x = Conv2D(32, kernel_size=[8, 8], padding='valid', strides=[4, 4], activation=None)(inputs)
        x = BatchNormalization(trainable=True, epsilon=1e-5)(x)
        x = Activation('relu')(x)
        x = Conv2D(64, kernel_size=[4, 4], padding='valid', strides=[2, 2], activation=None)(x)
        x = BatchNormalization(trainable=True, epsilon=1e-5)(x)
        x = Activation('relu')(x)
        x = Conv2D(64, kernel_size=(4, 4), padding='valid', strides=(2, 2), activation=None)(x)
        x = BatchNormalization(trainable=True, epsilon=1e-5)(x)
        x = Activation('relu')(x)
        x = Flatten()(x)
        x = Dense(512, activation='relu')(x)
        x = Dropout(0.5)(x)
        x = Dense(256, activation='relu')(x)
        x = Dropout(0.3)(x)
        x = Dense(128, activation='relu')(x)
        x = Dropout(0.05)(x)
        logits = Dense(self.action_space, activation=None)(x)

        self.model = tf.keras.Model(inputs=[inputs, rewards], outputs=logits)

        def policy_loss(r):
            def loss(labels, logits):
                policy = tf.nn.softmax(logits)
                entropy = tf.nn.softmax_cross_entropy_with_logits_v2(labels=policy, logits=logits)
                log = tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels, logits=logits)
                p_loss = log * tf.stop_gradient(r)
                p_loss = p_loss - 0.01 * entropy
                total_loss = tf.reduce_mean(p_loss)
                return total_loss

            return loss

        self.model.compile(optimizer=tf.train.AdamOptimizer(lr), loss=policy_loss(rewards))
        self.model.summary()

    def get_probs(self, s):
        s = s[np.newaxis, :]
        probs = self.model([s, np.array([1])]).numpy()
        probs = probs.squeeze()
        probs = self.softmax(probs)
        return probs

    def softmax(self, x):
        e_x = np.exp(x - np.max(x))
        return e_x / e_x.sum(axis=0)

    def update_policy(self, s, r, a):
        self.model.train_on_batch([s, r], a)
"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>
<span class="pl-k">from</span> <span class="pl-s1">tensorflow</span>.<span class="pl-s1">python</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span> <span class="pl-k">import</span> <span class="pl-c1">*</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-s1">tf</span>.<span class="pl-en">enable_eager_execution</span>()
<span class="pl-en">print</span>(<span class="pl-s1">tf</span>.<span class="pl-en">executing_eagerly</span>())


<span class="pl-k">class</span> <span class="pl-v">EagerTest</span>:
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">state_space</span>, <span class="pl-s1">action_space</span>, <span class="pl-s1">lr</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">action_space</span> <span class="pl-c1">=</span> <span class="pl-s1">action_space</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">lr</span> <span class="pl-c1">=</span> <span class="pl-s1">lr</span>

        <span class="pl-s1">inputs</span> <span class="pl-c1">=</span> <span class="pl-v">Input</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">84</span>, <span class="pl-c1">84</span>, <span class="pl-c1">4</span>))
        <span class="pl-s1">rewards</span> <span class="pl-c1">=</span> <span class="pl-v">Input</span>(<span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">1</span>,))

        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Conv2D</span>(<span class="pl-c1">32</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>[<span class="pl-c1">8</span>, <span class="pl-c1">8</span>], <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'valid'</span>, <span class="pl-s1">strides</span><span class="pl-c1">=</span>[<span class="pl-c1">4</span>, <span class="pl-c1">4</span>], <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">inputs</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Activation</span>(<span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Conv2D</span>(<span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>[<span class="pl-c1">4</span>, <span class="pl-c1">4</span>], <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'valid'</span>, <span class="pl-s1">strides</span><span class="pl-c1">=</span>[<span class="pl-c1">2</span>, <span class="pl-c1">2</span>], <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Activation</span>(<span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Conv2D</span>(<span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">4</span>, <span class="pl-c1">4</span>), <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-s">'valid'</span>, <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">2</span>, <span class="pl-c1">2</span>), <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Activation</span>(<span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Flatten</span>()(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-c1">512</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-c1">0.5</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-c1">256</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-c1">0.3</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-c1">128</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-c1">0.05</span>)(<span class="pl-s1">x</span>)
        <span class="pl-s1">logits</span> <span class="pl-c1">=</span> <span class="pl-v">Dense</span>(<span class="pl-s1">self</span>.<span class="pl-s1">action_space</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-c1">None</span>)(<span class="pl-s1">x</span>)

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-v">Model</span>(<span class="pl-s1">inputs</span><span class="pl-c1">=</span>[<span class="pl-s1">inputs</span>, <span class="pl-s1">rewards</span>], <span class="pl-s1">outputs</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)

        <span class="pl-k">def</span> <span class="pl-en">policy_loss</span>(<span class="pl-s1">r</span>):
            <span class="pl-k">def</span> <span class="pl-en">loss</span>(<span class="pl-s1">labels</span>, <span class="pl-s1">logits</span>):
                <span class="pl-s1">policy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax</span>(<span class="pl-s1">logits</span>)
                <span class="pl-s1">entropy</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax_cross_entropy_with_logits_v2</span>(<span class="pl-s1">labels</span><span class="pl-c1">=</span><span class="pl-s1">policy</span>, <span class="pl-s1">logits</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)
                <span class="pl-s1">log</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax_cross_entropy_with_logits_v2</span>(<span class="pl-s1">labels</span><span class="pl-c1">=</span><span class="pl-s1">labels</span>, <span class="pl-s1">logits</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)
                <span class="pl-s1">p_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">log</span> <span class="pl-c1">*</span> <span class="pl-s1">tf</span>.<span class="pl-en">stop_gradient</span>(<span class="pl-s1">r</span>)
                <span class="pl-s1">p_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">p_loss</span> <span class="pl-c1">-</span> <span class="pl-c1">0.01</span> <span class="pl-c1">*</span> <span class="pl-s1">entropy</span>
                <span class="pl-s1">total_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">reduce_mean</span>(<span class="pl-s1">p_loss</span>)
                <span class="pl-k">return</span> <span class="pl-s1">total_loss</span>

            <span class="pl-k">return</span> <span class="pl-s1">loss</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">compile</span>(<span class="pl-s1">optimizer</span><span class="pl-c1">=</span><span class="pl-s1">tf</span>.<span class="pl-s1">train</span>.<span class="pl-v">AdamOptimizer</span>(<span class="pl-s1">lr</span>), <span class="pl-s1">loss</span><span class="pl-c1">=</span><span class="pl-en">policy_loss</span>(<span class="pl-s1">rewards</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">summary</span>()

    <span class="pl-k">def</span> <span class="pl-en">get_probs</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>):
        <span class="pl-s1">s</span> <span class="pl-c1">=</span> <span class="pl-s1">s</span>[<span class="pl-s1">np</span>.<span class="pl-s1">newaxis</span>, :]
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">model</span>([<span class="pl-s1">s</span>, <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>])]).<span class="pl-en">numpy</span>()
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">probs</span>.<span class="pl-en">squeeze</span>()
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">softmax</span>(<span class="pl-s1">probs</span>)
        <span class="pl-k">return</span> <span class="pl-s1">probs</span>

    <span class="pl-k">def</span> <span class="pl-en">softmax</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-s1">e_x</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">exp</span>(<span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-s1">np</span>.<span class="pl-en">max</span>(<span class="pl-s1">x</span>))
        <span class="pl-k">return</span> <span class="pl-s1">e_x</span> <span class="pl-c1">/</span> <span class="pl-s1">e_x</span>.<span class="pl-en">sum</span>(<span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)

    <span class="pl-k">def</span> <span class="pl-en">update_policy</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>, <span class="pl-s1">r</span>, <span class="pl-s1">a</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">train_on_batch</span>([<span class="pl-s1">s</span>, <span class="pl-s1">r</span>], <span class="pl-s1">a</span>)</pre></div>
<p dir="auto"><strong>And same issue if I try to optimize manually</strong></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow as tf
from tensorflow.python.keras.layers import *

import numpy as np

tf.enable_eager_execution()
print(tf.executing_eagerly())


class EagerSeqTest:
    def __init__(self, state_space, action_space, lr):
        self.state_space = state_space
        self.action_space = action_space

        self.model = tf.keras.Sequential()

        self.model.add(InputLayer(input_shape=(84, 84, 4)))

        # Conv
        self.model.add(Conv2D(filters=32, kernel_size=(8, 8), strides=(4, 4), name='conv1'))
        self.model.add(BatchNormalization(trainable=True, epsilon=1e-5, name='batch_norm1'))
        self.model.add(ReLU(name='conv_1_out'))

        self.model.add(Conv2D(filters=64, kernel_size=(4, 4), strides=(2, 2), name='conv2'))
        self.model.add(BatchNormalization(trainable=True, epsilon=1e-5, name='batch_norm2'))
        self.model.add(ReLU(name='conv_2_out'))

        self.model.add(Conv2D(filters=64, kernel_size=(4, 4), strides=(2, 2), name='conv3'))
        self.model.add(BatchNormalization(trainable=True, epsilon=1e-5, name='batch_norm3'))
        self.model.add(ReLU(name='conv_3_out'))

        self.model.add(Flatten(name='flatten'))

        # Fully connected
        self.model.add(Dense(units=512, activation='relu', name='fc1'))
        self.model.add(Dropout(rate=0.4, name='dr1'))
        self.model.add(Dense(units=256, activation='relu', name='fc2'))
        self.model.add(Dropout(rate=0.3, name='dr2'))
        self.model.add(Dense(units=64, activation='relu', name='fc3'))
        self.model.add(Dropout(rate=0.03, name='dr3'))

        # Logits
        self.model.add(Dense(units=self.action_space, activation='linear', name='logits'))

        self.model.summary()

        # Optimizer
        self.optimizer = tf.train.AdamOptimizer(learning_rate=lr)

    def get_probs(self, s):
        s = s[np.newaxis, :]
        logits = self.model(s)
        probs = tf.nn.softmax(logits).numpy().squeeze()
        return probs

    def update_policy(self, s, r, a):
        with tf.GradientTape() as tape:
            loss = self.calc_loss(s, r, a)
        grads = tape.gradient(loss, self.model.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

    def calc_loss(self, s, r, a):
        logits = self.model(s)
        policy_loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels=a, logits=logits)
        policy_loss = tf.reduce_mean(policy_loss * tf.stop_gradient(r))
        return policy_loss"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>
<span class="pl-k">from</span> <span class="pl-s1">tensorflow</span>.<span class="pl-s1">python</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span> <span class="pl-k">import</span> <span class="pl-c1">*</span>

<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-s1">tf</span>.<span class="pl-en">enable_eager_execution</span>()
<span class="pl-en">print</span>(<span class="pl-s1">tf</span>.<span class="pl-en">executing_eagerly</span>())


<span class="pl-k">class</span> <span class="pl-v">EagerSeqTest</span>:
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">state_space</span>, <span class="pl-s1">action_space</span>, <span class="pl-s1">lr</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">state_space</span> <span class="pl-c1">=</span> <span class="pl-s1">state_space</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">action_space</span> <span class="pl-c1">=</span> <span class="pl-s1">action_space</span>

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-v">Sequential</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">InputLayer</span>(<span class="pl-s1">input_shape</span><span class="pl-c1">=</span>(<span class="pl-c1">84</span>, <span class="pl-c1">84</span>, <span class="pl-c1">4</span>)))

        <span class="pl-c"># Conv</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Conv2D</span>(<span class="pl-s1">filters</span><span class="pl-c1">=</span><span class="pl-c1">32</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">8</span>, <span class="pl-c1">8</span>), <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">4</span>, <span class="pl-c1">4</span>), <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'conv1'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'batch_norm1'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">ReLU</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'conv_1_out'</span>))

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Conv2D</span>(<span class="pl-s1">filters</span><span class="pl-c1">=</span><span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">4</span>, <span class="pl-c1">4</span>), <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">2</span>, <span class="pl-c1">2</span>), <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'conv2'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'batch_norm2'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">ReLU</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'conv_2_out'</span>))

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Conv2D</span>(<span class="pl-s1">filters</span><span class="pl-c1">=</span><span class="pl-c1">64</span>, <span class="pl-s1">kernel_size</span><span class="pl-c1">=</span>(<span class="pl-c1">4</span>, <span class="pl-c1">4</span>), <span class="pl-s1">strides</span><span class="pl-c1">=</span>(<span class="pl-c1">2</span>, <span class="pl-c1">2</span>), <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'conv3'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">BatchNormalization</span>(<span class="pl-s1">trainable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">epsilon</span><span class="pl-c1">=</span><span class="pl-c1">1e-5</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'batch_norm3'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">ReLU</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'conv_3_out'</span>))

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Flatten</span>(<span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'flatten'</span>))

        <span class="pl-c"># Fully connected</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dense</span>(<span class="pl-s1">units</span><span class="pl-c1">=</span><span class="pl-c1">512</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'fc1'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.4</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'dr1'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dense</span>(<span class="pl-s1">units</span><span class="pl-c1">=</span><span class="pl-c1">256</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'fc2'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.3</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'dr2'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dense</span>(<span class="pl-s1">units</span><span class="pl-c1">=</span><span class="pl-c1">64</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'relu'</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'fc3'</span>))
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dropout</span>(<span class="pl-s1">rate</span><span class="pl-c1">=</span><span class="pl-c1">0.03</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'dr3'</span>))

        <span class="pl-c"># Logits</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">add</span>(<span class="pl-v">Dense</span>(<span class="pl-s1">units</span><span class="pl-c1">=</span><span class="pl-s1">self</span>.<span class="pl-s1">action_space</span>, <span class="pl-s1">activation</span><span class="pl-c1">=</span><span class="pl-s">'linear'</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'logits'</span>))

        <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-en">summary</span>()

        <span class="pl-c"># Optimizer</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">train</span>.<span class="pl-v">AdamOptimizer</span>(<span class="pl-s1">learning_rate</span><span class="pl-c1">=</span><span class="pl-s1">lr</span>)

    <span class="pl-k">def</span> <span class="pl-en">get_probs</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>):
        <span class="pl-s1">s</span> <span class="pl-c1">=</span> <span class="pl-s1">s</span>[<span class="pl-s1">np</span>.<span class="pl-s1">newaxis</span>, :]
        <span class="pl-s1">logits</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">model</span>(<span class="pl-s1">s</span>)
        <span class="pl-s1">probs</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax</span>(<span class="pl-s1">logits</span>).<span class="pl-en">numpy</span>().<span class="pl-en">squeeze</span>()
        <span class="pl-k">return</span> <span class="pl-s1">probs</span>

    <span class="pl-k">def</span> <span class="pl-en">update_policy</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>, <span class="pl-s1">r</span>, <span class="pl-s1">a</span>):
        <span class="pl-k">with</span> <span class="pl-s1">tf</span>.<span class="pl-v">GradientTape</span>() <span class="pl-k">as</span> <span class="pl-s1">tape</span>:
            <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">calc_loss</span>(<span class="pl-s1">s</span>, <span class="pl-s1">r</span>, <span class="pl-s1">a</span>)
        <span class="pl-s1">grads</span> <span class="pl-c1">=</span> <span class="pl-s1">tape</span>.<span class="pl-en">gradient</span>(<span class="pl-s1">loss</span>, <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-s1">trainable_variables</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">optimizer</span>.<span class="pl-en">apply_gradients</span>(<span class="pl-en">zip</span>(<span class="pl-s1">grads</span>, <span class="pl-s1">self</span>.<span class="pl-s1">model</span>.<span class="pl-s1">trainable_variables</span>))

    <span class="pl-k">def</span> <span class="pl-en">calc_loss</span>(<span class="pl-s1">self</span>, <span class="pl-s1">s</span>, <span class="pl-s1">r</span>, <span class="pl-s1">a</span>):
        <span class="pl-s1">logits</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">model</span>(<span class="pl-s1">s</span>)
        <span class="pl-s1">policy_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softmax_cross_entropy_with_logits_v2</span>(<span class="pl-s1">labels</span><span class="pl-c1">=</span><span class="pl-s1">a</span>, <span class="pl-s1">logits</span><span class="pl-c1">=</span><span class="pl-s1">logits</span>)
        <span class="pl-s1">policy_loss</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-en">reduce_mean</span>(<span class="pl-s1">policy_loss</span> <span class="pl-c1">*</span> <span class="pl-s1">tf</span>.<span class="pl-en">stop_gradient</span>(<span class="pl-s1">r</span>))
        <span class="pl-k">return</span> <span class="pl-s1">policy_loss</span></pre></div>
<p dir="auto"><strong>Other info / logs</strong><br>
<strong>Keras outputs first episode</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[0.33398542 0.32970214 0.33631247]
[0.33452868 0.32947394 0.33599734]
[0.33389196 0.32968676 0.3364213 ]
[0.33383334 0.32978088 0.33638576]
[0.33387497 0.3296482  0.33647686]
[0.33390424 0.32952887 0.33656695]
[0.33359385 0.32956737 0.33683875]
[0.33371714 0.32963908 0.3366438 ]
[0.33354157 0.32960638 0.33685204]
[0.33366755 0.32962722 0.33670527]
            ..."><pre class="notranslate"><code class="notranslate">[0.33398542 0.32970214 0.33631247]
[0.33452868 0.32947394 0.33599734]
[0.33389196 0.32968676 0.3364213 ]
[0.33383334 0.32978088 0.33638576]
[0.33387497 0.3296482  0.33647686]
[0.33390424 0.32952887 0.33656695]
[0.33359385 0.32956737 0.33683875]
[0.33371714 0.32963908 0.3366438 ]
[0.33354157 0.32960638 0.33685204]
[0.33366755 0.32962722 0.33670527]
            ...
</code></pre></div>
<p dir="auto"><strong>Keras outputs after 10 episodes</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[0.2601601  0.31925505 0.42058486]
[0.3081141  0.35050547 0.34138042]
[0.30224514 0.41226208 0.28549278]
[0.2596298  0.4011653  0.33920488]
[0.24976756 0.44619036 0.30404213]
[0.22203408 0.50170064 0.2762653 ]
[0.25094017 0.460491   0.28856885]
[0.25819647 0.44147474 0.30032873]
[0.26891825 0.42985788 0.30122384]
[0.24224553 0.46432737 0.2934271 ]
              ..."><pre class="notranslate"><code class="notranslate">[0.2601601  0.31925505 0.42058486]
[0.3081141  0.35050547 0.34138042]
[0.30224514 0.41226208 0.28549278]
[0.2596298  0.4011653  0.33920488]
[0.24976756 0.44619036 0.30404213]
[0.22203408 0.50170064 0.2762653 ]
[0.25094017 0.460491   0.28856885]
[0.25819647 0.44147474 0.30032873]
[0.26891825 0.42985788 0.30122384]
[0.24224553 0.46432737 0.2934271 ]
              ...
</code></pre></div>
<p dir="auto"><strong>Eager outputs first episode</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[0.32854515 0.33827797 0.33317688]
[0.3286925  0.33824813 0.3330594 ]
[0.3283261  0.33858737 0.33308655]
[0.328541   0.3381491  0.33330992]
[0.328702   0.33778957 0.3335084 ]
[0.32902354 0.3377183  0.33325812]
[0.32862762 0.33784387 0.3335285 ]
[0.3286057  0.3380287  0.33336568]
[0.32879072 0.33825696 0.33295232]
[0.3287569  0.33825108 0.33299205]
              ..."><pre class="notranslate"><code class="notranslate">[0.32854515 0.33827797 0.33317688]
[0.3286925  0.33824813 0.3330594 ]
[0.3283261  0.33858737 0.33308655]
[0.328541   0.3381491  0.33330992]
[0.328702   0.33778957 0.3335084 ]
[0.32902354 0.3377183  0.33325812]
[0.32862762 0.33784387 0.3335285 ]
[0.3286057  0.3380287  0.33336568]
[0.32879072 0.33825696 0.33295232]
[0.3287569  0.33825108 0.33299205]
              ...
</code></pre></div>
<p dir="auto"><strong>Eager outputs after 10, 20, 100, .... episodes</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[0.3276771  0.33870533 0.33361754]
[0.3276953  0.33861196 0.33369276]
[0.3273576  0.33890665 0.33373576]
[0.32751966 0.3387702  0.33371013]
[0.3276457  0.33878204 0.33357224]
[0.32758534 0.3388276  0.33358708]
[0.3277567  0.3381556  0.33408773]
[0.32765684 0.33842626 0.33391687]
[0.32840943 0.3379219  0.33366865]
[0.32794353 0.33860624 0.33345023]
              ..."><pre class="notranslate"><code class="notranslate">[0.3276771  0.33870533 0.33361754]
[0.3276953  0.33861196 0.33369276]
[0.3273576  0.33890665 0.33373576]
[0.32751966 0.3387702  0.33371013]
[0.3276457  0.33878204 0.33357224]
[0.32758534 0.3388276  0.33358708]
[0.3277567  0.3381556  0.33408773]
[0.32765684 0.33842626 0.33391687]
[0.32840943 0.3379219  0.33366865]
[0.32794353 0.33860624 0.33345023]
              ...
</code></pre></div>
<p dir="auto"><em>Please make sure that this is a bug. As per our <a href="https://github.com/tensorflow/tensorflow/blob/master/ISSUES.md">GitHub Policy</a>, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template</em></p>
<p dir="auto"><strong>System information</strong></p>
<ul dir="auto">
<li>Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes</li>
<li>OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04.2 LTS (Bionic Beaver)</li>
<li>Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A</li>
<li>TensorFlow installed from (source or binary): source</li>
<li>TensorFlow version (use command below): 2.0.0-alpha0</li>
<li>Python version: 3.6.7</li>
<li>Bazel version (if compiling from source): N/A</li>
<li>GCC/Compiler version (if compiling from source): N/A</li>
<li>CUDA/cuDNN version: N/A</li>
<li>GPU model and memory: N/A</li>
</ul>
<p dir="auto">You can collect some of this information using our environment capture <a href="https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh">script</a><br>
You can also obtain the TensorFlow version with<br>
python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"</p>
<p dir="auto"><strong>Describe the current behavior</strong><br>
When training <code class="notranslate">tf.keras.fit(dataset)</code> with eager execution, where dataset is a <code class="notranslate">tf.data.Dataset</code>, I am finding a ~10x performance loss as compared to turning off eager execution.</p>
<p dir="auto"><strong>Describe the expected behavior</strong><br>
In the code attached below, I have tested training with and without eager execution, and with and without tf.data.Dataset.</p>
<p dir="auto">Here are the training times for one epoch of training on a 4 core CPU:<br>
Eager               + tf.data.Dataset : 219s  - 22ms/step<br>
Without Eager + tf.data.Dataset : 25s  - 3ms/step<br>
Eager               + numpy dataset : 26s  - 259us/sample<br>
Without Eager + numpy dataset : 26s  - 257us/sample</p>
<p dir="auto"><strong>Code to reproduce the issue</strong><br>
Provide a reproducible test case that is the bare minimum necessary to generate the problem.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import tensorflow as tf
import numpy as np

use_eager     = True
use_TFDataset = True

if not use_eager:
    tf.compat.v1.disable_eager_execution()

# Build dataset
n_data = 10**5
my_data = np.random.random((n_data,10,1))
my_targets = np.random.randint(0,2,(n_data,1))
data = ({'x_input':my_data}, {'target':my_targets})

# Create tf.data.Dataset
BATCH_SIZE = 10
dataset = tf.data.Dataset.from_tensor_slices(data)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(1)

#Build model
x_input = tf.keras.layers.Input((None,1), name='x_input')
RNN = tf.keras.layers.SimpleRNN(100, name='RNN')(x_input)
hidden = tf.keras.layers.Dense(100, name='hidden')(RNN)
dense = tf.keras.layers.Dense(1, name='target')(hidden)
my_model = tf.keras.models.Model(inputs = [x_input], outputs = [dense])
my_model.compile(optimizer='SGD', loss = 'binary_crossentropy')

# Train model
if use_TFDataset:
    my_model.fit(dataset, epochs = 1, steps_per_epoch=n_data//BATCH_SIZE) # divide by BATCH_SIZE to keep the number of training steps the same
else:
    my_model.fit(x = my_data, y = my_targets, epochs = 1, batch_size= BATCH_SIZE)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">tensorflow</span> <span class="pl-k">as</span> <span class="pl-s1">tf</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-s1">use_eager</span>     <span class="pl-c1">=</span> <span class="pl-c1">True</span>
<span class="pl-s1">use_TFDataset</span> <span class="pl-c1">=</span> <span class="pl-c1">True</span>

<span class="pl-k">if</span> <span class="pl-c1">not</span> <span class="pl-s1">use_eager</span>:
    <span class="pl-s1">tf</span>.<span class="pl-s1">compat</span>.<span class="pl-s1">v1</span>.<span class="pl-en">disable_eager_execution</span>()

<span class="pl-c"># Build dataset</span>
<span class="pl-s1">n_data</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span><span class="pl-c1">**</span><span class="pl-c1">5</span>
<span class="pl-s1">my_data</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-s1">random</span>.<span class="pl-en">random</span>((<span class="pl-s1">n_data</span>,<span class="pl-c1">10</span>,<span class="pl-c1">1</span>))
<span class="pl-s1">my_targets</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-s1">random</span>.<span class="pl-en">randint</span>(<span class="pl-c1">0</span>,<span class="pl-c1">2</span>,(<span class="pl-s1">n_data</span>,<span class="pl-c1">1</span>))
<span class="pl-s1">data</span> <span class="pl-c1">=</span> ({<span class="pl-s">'x_input'</span>:<span class="pl-s1">my_data</span>}, {<span class="pl-s">'target'</span>:<span class="pl-s1">my_targets</span>})

<span class="pl-c"># Create tf.data.Dataset</span>
<span class="pl-v">BATCH_SIZE</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>
<span class="pl-s1">dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">data</span>.<span class="pl-v">Dataset</span>.<span class="pl-en">from_tensor_slices</span>(<span class="pl-s1">data</span>)
<span class="pl-s1">dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">dataset</span>.<span class="pl-en">batch</span>(<span class="pl-v">BATCH_SIZE</span>)
<span class="pl-s1">dataset</span> <span class="pl-c1">=</span> <span class="pl-s1">dataset</span>.<span class="pl-en">prefetch</span>(<span class="pl-c1">1</span>)

<span class="pl-c">#Build model</span>
<span class="pl-s1">x_input</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Input</span>((<span class="pl-c1">None</span>,<span class="pl-c1">1</span>), <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'x_input'</span>)
<span class="pl-v">RNN</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">SimpleRNN</span>(<span class="pl-c1">100</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'RNN'</span>)(<span class="pl-s1">x_input</span>)
<span class="pl-s1">hidden</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-c1">100</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'hidden'</span>)(<span class="pl-v">RNN</span>)
<span class="pl-s1">dense</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">layers</span>.<span class="pl-v">Dense</span>(<span class="pl-c1">1</span>, <span class="pl-s1">name</span><span class="pl-c1">=</span><span class="pl-s">'target'</span>)(<span class="pl-s1">hidden</span>)
<span class="pl-s1">my_model</span> <span class="pl-c1">=</span> <span class="pl-s1">tf</span>.<span class="pl-s1">keras</span>.<span class="pl-s1">models</span>.<span class="pl-v">Model</span>(<span class="pl-s1">inputs</span> <span class="pl-c1">=</span> [<span class="pl-s1">x_input</span>], <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> [<span class="pl-s1">dense</span>])
<span class="pl-s1">my_model</span>.<span class="pl-en">compile</span>(<span class="pl-s1">optimizer</span><span class="pl-c1">=</span><span class="pl-s">'SGD'</span>, <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-s">'binary_crossentropy'</span>)

<span class="pl-c"># Train model</span>
<span class="pl-k">if</span> <span class="pl-s1">use_TFDataset</span>:
    <span class="pl-s1">my_model</span>.<span class="pl-en">fit</span>(<span class="pl-s1">dataset</span>, <span class="pl-s1">epochs</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span>, <span class="pl-s1">steps_per_epoch</span><span class="pl-c1">=</span><span class="pl-s1">n_data</span><span class="pl-c1">//</span><span class="pl-v">BATCH_SIZE</span>) <span class="pl-c"># divide by BATCH_SIZE to keep the number of training steps the same</span>
<span class="pl-k">else</span>:
    <span class="pl-s1">my_model</span>.<span class="pl-en">fit</span>(<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">my_data</span>, <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">my_targets</span>, <span class="pl-s1">epochs</span> <span class="pl-c1">=</span> <span class="pl-c1">1</span>, <span class="pl-s1">batch_size</span><span class="pl-c1">=</span> <span class="pl-v">BATCH_SIZE</span>)</pre></div>
<p dir="auto"><strong>Other info / logs</strong><br>
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.</p>