tensorflow2 自定义损失函数使用的隐藏坑

import numpy as np
import tensorflow as tf
from tensorflow import keras
x = np.random.random((1000, 32))
y = np.random.random((1000, 1))
class CustomModel(keras.Model):
 tf.random.set_seed(100)
 def train_step(self, data):
  # Unpack the data. Its structure depends on your model and
  # on what you pass to `fit()`.
  x, y = data
  with tf.GradientTape() as tape:
y_pred = self(x, training=True)  # Forward pass
# Compute the loss value
# (the loss function is configured in `compile()`)
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
  return {m.name: m.result() for m in self.metrics}
 

# Construct and compile an instance of CustomModel
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = CustomModel(inputs, outputs)
model.compile(optimizer="adam", loss=tf.losses.MSE, metrics=["mae"])
# Just use `fit` as usual
model.fit(x, y, epochs=1, shuffle=False)
32/32 [==============================] - 0s 1ms/step - loss: 0.2783 - mae: 0.4257
 
<tensorflow.python.keras.callbacks.History at 0x7ff7edf6dfd0>

def custom_mse(y_true, y_pred):
 return tf.reduce_mean((y_true - y_pred)**2, axis=-1)
a_true = tf.constant([1., 1.5, 1.2])
a_pred = tf.constant([1., 2, 1.5])
custom_mse(a_true, a_pred)
<tf.Tensor: shape=(), dtype=float32, numpy=0.11333332>
tf.losses.MSE(a_true, a_pred)
<tf.Tensor: shape=(), dtype=float32, numpy=0.11333332>

my_model = CustomModel(inputs, outputs)
my_model.compile(optimizer="adam", loss=custom_mse, metrics=["mae"])
my_model.fit(x, y, epochs=1, shuffle=False)
32/32 [==============================] - 0s 820us/step - loss: 0.1628 - mae: 0.3257
<tensorflow.python.keras.callbacks.History at 0x7ff7edeb7810>

loss_tracker = keras.metrics.Mean(name="loss")
mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
class MyCustomModel(keras.Model):
 tf.random.set_seed(100)
 def train_step(self, data):
  # Unpack the data. Its structure depends on your model and
  # on what you pass to `fit()`.
  x, y = data
  with tf.GradientTape() as tape:
y_pred = self(x, training=True)  # Forward pass
# Compute the loss value
# (the loss function is configured in `compile()`)
loss = custom_mse(y, y_pred)
# loss += 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))
  
  # 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()}
 
 @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, mae_metric]
 
# Construct and compile an instance of CustomModel
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
my_model_beta = MyCustomModel(inputs, outputs)
my_model_beta.compile(optimizer="adam")
# Just use `fit` as usual
my_model_beta.fit(x, y, epochs=1, shuffle=False)
32/32 [==============================] - 0s 960us/step - loss: 0.2783 - mae: 0.4257
<tensorflow.python.keras.callbacks.History at 0x7ff7eda3d810>