Keras自定义损失函数中使用TensorFlow会话：基于微分差异拟合二元函数

我最近折腾了个有意思的练习：用Keras拟合二元函数 ( f(x,y)=x^2+y2 )，但损失函数不是常规的预测值和真实值的偏差，而是解析微分与自动微分的差值的平方和。过程中参考了Stack Overflow上两个核心问题的思路，结合Keras文档完成了代码，分享给大家：

核心问题拆解 #

要实现这个需求，得解决两个关键问题：

• 怎么在Keras的自定义损失函数里访问当前的输入样本
• 怎么获取模型输出相对于输入的自动微分结果

我参考的两个思路方向：

• 自定义损失函数中访问输入样本的实现方式
• 获取模型输出对输入（而非权重）的梯度方法

完整代码实现 #

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense
import numpy as np

# 1. 定义目标二元函数 f(x,y) = x² + y²
def target_func(x):
    return tf.reduce_sum(tf.square(x), axis=1)

# 2. 定义该函数的解析梯度：df/dx=2x，df/dy=2y
def analytical_grad(x):
    return 2 * x

# 3. 自定义损失函数：计算自动微分与解析微分的平方差损失
def grad_based_loss(y_true, y_pred, input_data):
    # 使用GradientTape记录计算过程，获取自动微分
    with tf.GradientTape() as tape:
        tape.watch(input_data)
        pred = model(input_data, training=True)
    auto_grad = tape.gradient(pred, input_data)
    
    # 计算两种梯度的差值的平方均值作为损失
    grad_diff = auto_grad - analytical_grad(input_data)
    return tf.reduce_mean(tf.square(grad_diff))

# 4. 构建神经网络模型
input_layer = Input(shape=(2,))
hidden = Dense(16, activation='relu')(input_layer)
output = Dense(1)(hidden)
model = Model(inputs=input_layer, outputs=output)

# 5. 包装损失函数，让它能接收输入样本
def loss_wrapper(input_tensor):
    def loss(y_true, y_pred):
        return grad_based_loss(y_true, y_pred, input_tensor)
    return loss

# 6. 编译并训练模型
model.compile(optimizer='adam', loss=loss_wrapper(input_layer))

# 生成训练数据：[-5,5]范围内的随机二元样本
x_train = np.random.uniform(-5, 5, (1000, 2))
y_train = target_func(x_train).reshape(-1, 1)

# 开始训练
model.fit(x_train, y_train, epochs=50, batch_size=32)

关键细节说明 #

• 梯度计算：用tf.GradientTape来追踪输入的计算过程，从而得到模型输出对输入的自动微分，这是TensorFlow自动微分的核心用法
• 损失函数包装：因为Keras默认的损失函数只接收y_true和y_pred，所以我们需要用一个包装函数把输入层传入损失函数，让它能访问到当前的输入样本
• 训练逻辑：虽然损失用的是梯度差值，但模型最终还是会拟合目标函数，因为只有当模型的输出足够接近目标函数时，它的梯度才会和解析梯度一致

内容的提问来源于stack exchange，提问作者Lucas Farias