Issue 1:
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.

Issue 2:
#

Click to expand!

Issue Type

Bug

Source

binary

Tensorflow Version

2.8.2

Custom Code

Yes

OS Platform and Distribution

Linux Ubuntu 18.04.5

Python version

Python 3.7.13

Current Behavior

I am following the steps here for customizing what happens in train_step when you customize a model:

https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit

The issue is, I am getting very different results for the loss when I am using the exact same tensorflow.keras.losses.mean_squared_error function in different locations.

When I use mean_squared_error "manually" during train_step (I am calling this the non-compiled mse), the loss is significantly worse than if I pass it into the compile method, e.g. model.compile(optimizer="adam", loss=mean_squared_error) (I am calling this the compiled mse).

tensorflow.keras.metrics.MeanSquaredError yields roughly the same result as the compiled version of mse, but it is still not exactly the same.

Expected Behavior

My expectation is that the loss should be the same whether it is passed into .compile, defined during train_step, or used as a metric. But it is significantly worse if I do not pass it into .compile and I don't know why. I am following the documentation, and as far as I can tell am not doing anything incorrectly.

Why is the loss so different when it is and is not passed into .compile?

Colab Notebook

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.

https://colab.research.google.com/drive/1c1L8KJYQbAldphKvFvmw5cHi9EyBQP0X?usp=sharing

Originally I opened a ticket at tensorflow/tensorflow here, and they were able to reproduce the issue. This is the gist:

https://colab.research.google.com/gist/sushreebarsa/f1144243be9eed060c39c60611811811/56866.ipynb

Standalone code to reproduce behavior

This is the same as what is in the Colab notebooks.

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()

Related Issues

I could not find an issue that was the same as this one, but 16834 might be related, as it also notes a difference between compiled mean squared error and manually calculated mean squared error:

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.

I have not tested whether or not this behavior is reproduced when using a loss other than mean squared error.