具有自定义渐变的自定义训练损失

Question

我正尝试在Tensorflow v2中编写自定义损失，为简单起见，我说我正在使用均方误差损失，如下所示，

loss_object = tf.keras.losses.MeanSquaredError()


def loss(model, x, y, training):
  # training=training is needed only if there are layers with different
  # behavior during training versus inference (e.g. Dropout).
  y_ = model(x, training=training)
  return loss_object(y_true=y, y_pred=y_)

现在我知道Tensorflow确实automatic differentiation。

但是我想在BackPropagation算法中指定我的自定义渐变， 如果我们使用MSE，则必须执行以下

是否有可能在 Keras 中将替换为，其中p是在训练期间通过渐变之前应用的张量。