Model.evaluate returns 0 loss when using custom model

Question

I am trying to use my own train step in with Keras by creating a class that inherits from Model. It seems that the training works correctly but the evaluate function always returns 0 on the loss even if I send to it the train data, which have a big loss value during the training. I can't share my code but was able to reproduce using the example form the Keras api in https://keras.io/guides/customizing_what_happens_in_fit/ I changed the Dense layer to have 2 units instead of one, and made its activation to sigmoid.

The code:

import tensorflow.keras as keras
import tensorflow as tf
import numpy as np

loss_tracker = keras.metrics.Mean(name="loss")
# mae_metric = keras.metrics.MeanAbsoluteError(name="mae")

class CustomModel(keras.Model):
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute our own loss
            loss = self.compiled_loss(y, y_pred)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # Compute our own metrics
        loss_tracker.update_state(loss)
        self.compiled_metrics.update_state(y, y_pred)
        return {"loss": loss_tracker.result(), "mae": self.compiled_metrics.metrics[0].result()}

    @property
    def metrics(self):
        # We list our `Metric` objects here so that `reset_states()` can be
        # called automatically at the start of each epoch
        # or at the start of `evaluate()`.
        # If you don't implement this property, you have to call
        # `reset_states()` yourself at the time of your choosing.
        return [loss_tracker] + self.compiled_metrics.metrics


if __name__ == '__main__':
    # Construct an instance of CustomModel
    inputs = keras.Input(shape=(32,))
    outputs = keras.layers.Dense(2, activation='sigmoid')(inputs)
    model = CustomModel(inputs, outputs)

    # We don't passs a loss or metrics here.
    model.compile(optimizer="adam", loss='mean_squared_error', metrics=['mean_squared_error'])

    # Just use `fit` as usual -- you can use callbacks, etc.
    x = np.random.random((1000, 32))
    y = np.random.random((1000, 1))
    model.fit(x, y, epochs=5)
    print(model.evaluate(x, y, return_dict=True))
    x = np.random.random((1000, 32))
    y = np.random.random((1000, 1))
    print(model.evaluate(x, y, return_dict=True))

I try to run evaluate on the original train data as well on some random data and both return 0 on the loss and MAE.

The output: Epoch 1/5<c 32/32 [==============================] - 0s 708us/step - loss: 0.1133 - mae: 0.1193
Epoch 2/5
32/32 [==============================] - 0s 688us/step - loss: 0.0962 - mae: 0.1000
Epoch 3/5
32/32 [==============================] - 0s 677us/step - loss: 0.0926 - mae: 0.0915
Epoch 4/5
32/32 [==============================] - 0s 740us/step - loss: 0.0922 - mae: 0.0872
Epoch 5/5
32/32 [==============================] - 0s 708us/step - loss: 0.0912 - mae: 0.0938
Evaluations:
32/32 [==============================] - 0s 1ms/step - loss: 0.0000e+00 - mean_squared_error: 0.0916
{'loss': 0.0, 'mean_squared_error': 0.09163134545087814}
32/32 [==============================] - 0s 1ms/step - loss: 0.0000e+00 - mean_squared_error: 0.0894
{'loss': 0.0, 'mean_squared_error': 0.08940737694501877}

The first evaluation is on the training data and returns 0 although in the training it isn't 0 and also the random data returns 0 - no way it is really 0.

Any ideas what might cause this problem? Or how to solve this except for overriding myself the eval method?

Edit: After M.Innat response I changed the model to use the compiled loss and metrics now the metrics work, the compiled loss still doesn't.

Answer 1

As you manually use the loss and metrics function in the train_step (not in the .compile) for the training set, you should also do the same for the validation set or by defining the test_step in the custom model in order to get the loss score and metrics score. Add the following function to your custom model.

    def test_step(self, data):
        # Unpack the data
        x, y = data
        # Compute predictions
        y_pred = self(x, training=False)
        loss = keras.losses.mean_squared_error(y, y_pred)
        
        # Compute our own metrics
        loss_tracker.update_state(loss)
        mae_metric.update_state(y, y_pred)
        return {"loss": loss_tracker.result(), "mae": mae_metric.result()}

Answer 2

I found what caused the problem in the example. It is the usage of a custom loss tracker and overriding the metrics function.
Actually if you use compiled loss Keras will track it by itself in the self.compiled_loss.metrics.

The problem was when using a custom loss I had to override the metrics function to get a correct result and then the test_step didn't use collect the loss correctly.

The code that does work without overriding the test_step function:

import tensorflow.keras as keras
import tensorflow as tf

class CustomModel(keras.Model):
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute our own loss
            loss = self.compiled_loss(y, y_pred)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # Compute our own metrics
        self.compiled_metrics.update_state(y, y_pred)
        return {"loss": self.compiled_loss.metrics[0].result(), "mae": self.compiled_metrics.metrics[0].result()}