<p dir="auto">When using keras for regression problems, loss uses mean squared error. The loss displayed by the network fit in the last epoch during training is inconsistent with the loss calculated from the input data after the model is saved.</p>
<p dir="auto">#</p><details><summary>Click to expand!</summary><p dir="auto"></p>
<h3 dir="auto">Issue Type</h3>
<p dir="auto">Bug</p>
<h4 dir="auto">Source</h4>
<p dir="auto">binary</p>
<h4 dir="auto">Tensorflow Version</h4>
<p dir="auto">2.8.2</p>
<h4 dir="auto">Custom Code</h4>
<p dir="auto">Yes</p>
<h4 dir="auto">OS Platform and Distribution</h4>
<p dir="auto">Linux Ubuntu 18.04.5</p>
<h4 dir="auto">Python version</h4>
<p dir="auto">Python 3.7.13</p>
<h4 dir="auto">Current Behavior</h4>
<p dir="auto">I am following the steps here for customizing what happens in <code class="notranslate">train_step</code> when you customize a model:</p>
<p dir="auto"><a href="https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit" rel="nofollow">https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit</a></p>
<p dir="auto">The issue is, I am getting very different results for the loss when I am using the exact same <code class="notranslate">tensorflow.keras.losses.mean_squared_error</code> function in different locations.</p>
<p dir="auto">When I use <code class="notranslate">mean_squared_error</code> "manually" during <code class="notranslate">train_step</code> (I am calling this the non-compiled mse), the loss is significantly worse than if I  pass it into the compile method, e.g. <code class="notranslate">model.compile(optimizer="adam", loss=mean_squared_error)</code> (I am calling this the compiled mse).</p>
<p dir="auto"><code class="notranslate">tensorflow.keras.metrics.MeanSquaredError</code> yields roughly the same result as the compiled version of mse, but it is still not exactly the same.</p>
<h4 dir="auto">Expected Behavior</h4>
<p dir="auto">My expectation is that the loss should be the same whether it is passed into <code class="notranslate">.compile</code>, defined during <code class="notranslate">train_step</code>, or used as a metric. But it is significantly worse if I do not pass it into <code class="notranslate">.compile</code> and I don't know why. I am following the documentation, and as far as I can tell am not doing anything incorrectly.</p>
<p dir="auto">Why is the loss so different when it is and is not passed into <code class="notranslate">.compile</code>?</p>
<h4 dir="auto">Colab Notebook</h4>
<p dir="auto">I have reproduced this behavior in this google colab notebook, link below. The plot shows that non-compiled mse is systematically worse than compiled mse.</p>
<p dir="auto"><a href="https://colab.research.google.com/drive/1c1L8KJYQbAldphKvFvmw5cHi9EyBQP0X?usp=sharing" rel="nofollow">https://colab.research.google.com/drive/1c1L8KJYQbAldphKvFvmw5cHi9EyBQP0X?usp=sharing</a></p>
<p dir="auto">Originally I opened a ticket at tensorflow/tensorflow <a href="https://github.com/tensorflow/tensorflow/issues/56866" data-hovercard-type="issue" data-hovercard-url="/tensorflow/tensorflow/issues/56866/hovercard">here</a>, and they were able to reproduce the issue. This is the gist:</p>
<p dir="auto"><a href="https://colab.research.google.com/gist/sushreebarsa/f1144243be9eed060c39c60611811811/56866.ipynb" rel="nofollow">https://colab.research.google.com/gist/sushreebarsa/f1144243be9eed060c39c60611811811/56866.ipynb</a></p>
<h4 dir="auto">Standalone code to reproduce behavior</h4>
<p dir="auto">This is the same as what is in the Colab notebooks.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import tensorflow as tf
from tensorflow.keras.metrics import Mean, MeanSquaredError
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Dropout, Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import mean_squared_error
import matplotlib.pyplot as plt


seed = 113
epochs = 30
batch_size = 10

n = 5
history_dict = {}

# load dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(
    path=&quot;boston_housing.npz&quot;, test_split=0.2, seed=seed,
)
n_validate = int(0.1 * len(x_train))
train_dataset = tf.data.Dataset.from_tensor_slices(
    (x_train[:-n_validate], y_train[:-n_validate])
).batch(batch_size)
valid_dataset = tf.data.Dataset.from_tensor_slices(
    (x_train[-n_validate:], y_train[-n_validate:])
).batch(batch_size)
test_dataset = tf.data.Dataset.from_tensor_slices(
    (x_test, y_test)
).batch(batch_size)

for i in range(n):
    
    tf.keras.backend.clear_session()
    
    mse_tracker = Mean(name=&quot;non_compiled_mse&quot;)
    mse_tracker_metric = MeanSquaredError(name=&quot;mse_metric&quot;)
    
    class CustomModel(Model):
        
        @tf.function
        def train_step(self, data):
            x, y_true = data
            with tf.GradientTape() as tape:
                y_pred = self(x, training=True)
                compiled_mse = self.compiled_loss(y_true, y_pred)
                non_compiled_mse = mean_squared_error(y_true, y_pred)
                mse_tracker.update_state(non_compiled_mse)
                mse_tracker_metric.update_state(y_true, y_pred)
            gradients = tape.gradient(compiled_mse, self.trainable_variables)
            self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
            return {m.name: m.result() for m in self.metrics}
        
        @tf.function()
        def test_step(self, data):
            x, y_true = data
            y_pred = self(x, training=False)
            compiled_mse = self.compiled_loss(y_true, y_pred)
            non_compiled_mse = mean_squared_error(y_true, y_pred)
            mse_tracker.update_state(non_compiled_mse)
            mse_tracker_metric.update_state(y_true, y_pred)
            return {m.name: m.result() for m in self.metrics}

        @property
        def metrics(self):
            return super().metrics + [mse_tracker, mse_tracker_metric]
    
    input_layer = Input(shape=x_train.shape[1:])
    x = Dense(8, activation='linear')(input_layer)
    x = Dense(8, activation='linear')(x)
    output_layer = Dense(units=1, activation='linear')(x)
    model = CustomModel(input_layer, output_layer)
    model.compile(optimizer=Adam(learning_rate=1e-3), loss=mean_squared_error)
    history = model.fit(
        train_dataset,
        validation_data=valid_dataset,
        epochs=epochs,
        verbose=0,
    )
    history_dict[f&quot;compiled_mse_{i}&quot;] = history.history[&quot;loss&quot;]
    history_dict[f&quot;non_compiled_mse_{i}&quot;] = history.history[&quot;non_compiled_mse&quot;]
    history_dict[f&quot;mse_metric_{i}&quot;] = history.history[&quot;mse_metric&quot;]


plot_start_epoch = 5
plt.figure(figsize=(12, 6))
for i in range(n):
    plt.plot(
        range(epochs)[plot_start_epoch:],
        history_dict[f&quot;compiled_mse_{i}&quot;][plot_start_epoch:],
        label=f&quot;compiled_mse&quot; if i == 0 else None,
        color=&quot;black&quot;,
    )
    plt.plot(
        range(epochs)[plot_start_epoch:],
        history_dict[f&quot;non_compiled_mse_{i}&quot;][plot_start_epoch:],
        label=f&quot;non_compiled_mse&quot; if i == 0 else None,
        color=&quot;orange&quot;,
    )
    plt.plot(
        range(epochs)[plot_start_epoch:],
        history_dict[f&quot;mse_metric_{i}&quot;][plot_start_epoch:],
        label=f&quot;mse_metric&quot; if i == 0 else None,
        color=&quot;red&quot;,
        linestyle=&quot;dashed&quot;,
    )
plt.ylabel(&quot;loss&quot;)
plt.xlabel(&quot;epoch&quot;)
plt.title(&quot;compiled vs non-compiled mse vs mse metric&quot;)
plt.legend()
plt.show()"><pre class="notranslate"><code class="notranslate">import tensorflow as tf
from tensorflow.keras.metrics import Mean, MeanSquaredError
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Dropout, Input
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.losses import mean_squared_error
import matplotlib.pyplot as plt


seed = 113
epochs = 30
batch_size = 10

n = 5
history_dict = {}

# load dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(
    path="boston_housing.npz", test_split=0.2, seed=seed,
)
n_validate = int(0.1 * len(x_train))
train_dataset = tf.data.Dataset.from_tensor_slices(
    (x_train[:-n_validate], y_train[:-n_validate])
).batch(batch_size)
valid_dataset = tf.data.Dataset.from_tensor_slices(
    (x_train[-n_validate:], y_train[-n_validate:])
).batch(batch_size)
test_dataset = tf.data.Dataset.from_tensor_slices(
    (x_test, y_test)
).batch(batch_size)

for i in range(n):
    
    tf.keras.backend.clear_session()
    
    mse_tracker = Mean(name="non_compiled_mse")
    mse_tracker_metric = MeanSquaredError(name="mse_metric")
    
    class CustomModel(Model):
        
        @tf.function
        def train_step(self, data):
            x, y_true = data
            with tf.GradientTape() as tape:
                y_pred = self(x, training=True)
                compiled_mse = self.compiled_loss(y_true, y_pred)
                non_compiled_mse = mean_squared_error(y_true, y_pred)
                mse_tracker.update_state(non_compiled_mse)
                mse_tracker_metric.update_state(y_true, y_pred)
            gradients = tape.gradient(compiled_mse, self.trainable_variables)
            self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
            return {m.name: m.result() for m in self.metrics}
        
        @tf.function()
        def test_step(self, data):
            x, y_true = data
            y_pred = self(x, training=False)
            compiled_mse = self.compiled_loss(y_true, y_pred)
            non_compiled_mse = mean_squared_error(y_true, y_pred)
            mse_tracker.update_state(non_compiled_mse)
            mse_tracker_metric.update_state(y_true, y_pred)
            return {m.name: m.result() for m in self.metrics}

        @property
        def metrics(self):
            return super().metrics + [mse_tracker, mse_tracker_metric]
    
    input_layer = Input(shape=x_train.shape[1:])
    x = Dense(8, activation='linear')(input_layer)
    x = Dense(8, activation='linear')(x)
    output_layer = Dense(units=1, activation='linear')(x)
    model = CustomModel(input_layer, output_layer)
    model.compile(optimizer=Adam(learning_rate=1e-3), loss=mean_squared_error)
    history = model.fit(
        train_dataset,
        validation_data=valid_dataset,
        epochs=epochs,
        verbose=0,
    )
    history_dict[f"compiled_mse_{i}"] = history.history["loss"]
    history_dict[f"non_compiled_mse_{i}"] = history.history["non_compiled_mse"]
    history_dict[f"mse_metric_{i}"] = history.history["mse_metric"]


plot_start_epoch = 5
plt.figure(figsize=(12, 6))
for i in range(n):
    plt.plot(
        range(epochs)[plot_start_epoch:],
        history_dict[f"compiled_mse_{i}"][plot_start_epoch:],
        label=f"compiled_mse" if i == 0 else None,
        color="black",
    )
    plt.plot(
        range(epochs)[plot_start_epoch:],
        history_dict[f"non_compiled_mse_{i}"][plot_start_epoch:],
        label=f"non_compiled_mse" if i == 0 else None,
        color="orange",
    )
    plt.plot(
        range(epochs)[plot_start_epoch:],
        history_dict[f"mse_metric_{i}"][plot_start_epoch:],
        label=f"mse_metric" if i == 0 else None,
        color="red",
        linestyle="dashed",
    )
plt.ylabel("loss")
plt.xlabel("epoch")
plt.title("compiled vs non-compiled mse vs mse metric")
plt.legend()
plt.show()
</code></pre></div>
<h4 dir="auto">Related Issues</h4>
<p dir="auto">I could not find an issue that was the same as this one, but <a href="https://github.com/keras-team/keras/issues/16834" data-hovercard-type="issue" data-hovercard-url="/keras-team/keras/issues/16834/hovercard">16834</a> might be related, as it also notes a difference between compiled mean squared error and manually calculated mean squared error:</p>
<blockquote>
<p dir="auto">When using keras for regression problems, loss uses mean squared error. The loss displayed by the network fit in the last epoch during training is inconsistent with the loss calculated from the input data after the model is saved.</p>
</blockquote>
<p dir="auto">I have not tested whether or not this behavior is reproduced when using a loss other than mean squared error.</p></details><p dir="auto"></p>