Evaluate Tensorflow Tensor inside train_step() when run_eagerly = False

Question

I'm running into issues trying to evaluate Tensors computed inside a custom train_step() when I subclass the Model class. I can use Tensor.numpy() when I pass run_eagerly = True inside model.compile(...) but to my understanding this is not efficient. I have tried other suggestions I found like using Tensor.eval() (no default session), backend.get_values(Tensor) (Tensor has no attribute numpy()) etc. with no success. I have a simplified example of what I would like to achieve below:

class CustomModel(keras.Model):

def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self.saved_pred = []

def train_step(self, data):
    x, y = data
    
    with tf.GradientTape() as tape:
        y_pred = self(x, training=True)  # Forward pass
        loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)

    # Compute gradients
    trainable_vars = self.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    # Update weights
    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    # Update metrics (includes the metric that tracks the loss)
    self.compiled_metrics.update_state(y, y_pred)
    # Return a dict mapping metric names to current value

    self.saved_pred.append(y_pred)

    return {m.name: m.result() for m in self.metrics}

import numpy as np

inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs = inputs, outputs = outputs)
model.compile(optimizer=“adam”, loss=“mse”, metrics=[“mae”])

x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
model.fit(x, y, epochs=3)

When I print model.saved_pred, I get the following:

ListWrapper([, 
])

Is there some way to extract the values of these Tensors (as numpy arrays), either inside train_step() or after model.fit()?

user_9551404 · Answer

From the keras doc:

run eagerly: Bool. Defaults to False . If True , this Model 's logic will not be wrapped in a tf.function. Recommended to leave this as None unless your Model cannot be run inside a tf.function...

what it does (copied from Sayak_Paul's answer):

When something is wrapped inside tf.function it has the advantage of being run in graph mode. All the backend compilation engineering is handled by TensorFlow itself in this case. The advantage is when all the operations are available as a graph we know much resources to allocate and how to best optimize the graph with the available resources. For more details refer to the following: https://www.tensorflow.org/guide/function

run_eagerly=True lets figure out what exactly is going inside your model training loop. Let’s say you have implemented a custom loop and put that inside the train_step() method of a subclasses model. Setting run_eagerly to True will help you debug that loop if anything goes wrong. For practical applications of this, refer to the following guide: https://keras.io/examples/keras_recipes/debugging_tips/

Therefore, it's not efficient, but let's you debug your model. Because it will be awfully slow, set the run_eagerly flag to True as soon as you start training.

Evaluate Tensorflow Tensor inside train_step() when run_eagerly = False

1 Answers