如何在tensorflow中实现自定义损失函数和梯度反向传播

import tensorflow as tf

# 定义自定义损失函数
def custom_loss(y_true, y_pred):
    # 自定义损失计算逻辑
    loss = tf.square(y_true - y_pred)
    return loss

# 创建模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 定义输入和标签
inputs = tf.random.normal([32, 10])
labels = tf.random.normal([32, 10])

# 计算损失
with tf.GradientTape() as tape:
    predictions = model(inputs)
    loss = tf.reduce_mean(custom_loss(labels, predictions))

# 计算梯度
gradients = tape.gradient(loss, model.trainable_variables)

# 使用优化器更新模型参数
optimizer = tf.keras.optimizers.Adam()
optimizer.apply_gradients(zip(gradients, model.trainable_variables))