Keras的核心原則是逐步揭示復雜性，可以在保持相應的高級便利性的同時，對操作細節進行更多控制。當我們要自定義fit中的訓練算法時，可以重寫模型中的train_step方法，然后調用fit來訓練模型。

這里以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 CustomModelinputs = 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 usualmodel.fit(x, y, epochs=1, shuffle=False)
32/32 [==============================] - 0s 1ms/step - loss: 0.2783 - mae: 0.4257

這里的loss是tensorflow庫中實現了的損失函數，如果想自定義損失函數，然后將損失函數傳入model.compile中，能正常按我們預想的work嗎？

答案竟然是否定的，而且沒有錯誤提示，只是loss計算不會符合我們的預期。

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.losses.MSE(a_true, a_pred)

以上結果證實了我們自定義loss的正確性，下面我們直接將自定義的loss置入compile中的loss參數中，看看會發生什么。

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

我們看到，這里的loss與我們與標準的tf.losses.MSE明顯不同。這說明我們自定義的loss以這種方式直接傳遞進model.compile中，是完全錯誤的操作。

正確運用自定義loss的姿勢是什么呢？下面揭曉。

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 CustomModelinputs = 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 usualmy_model_beta.fit(x, y, epochs=1, shuffle=False)
32/32 [==============================] - 0s 960us/step - loss: 0.2783 - mae: 0.4257

終于，通過跳過在 compile() 中傳遞損失函數，而在 train_step 中手動完成所有計算內容，我們獲得了與之前默認tf.losses.MSE完全一致的輸出，這才是我們想要的結果。