TF2.6下Keras用GradientTape实现自定义损失函数报错如何解决

报错根因 #

报错的核心原因是Keras模型编译阶段，model.input是静态的KerasTensor占位符对象，不是模型运行时实际喂入的tf.Tensor类型输入，无法直接通过tf.convert_to_tensor转换；同时在损失函数内部调用model(x_tensor)会重复构建前向计算逻辑，引发计算图嵌套冲突。

推荐解决方案：使用自定义训练循环 #

你要实现的物理信息神经网络（PINN）损失逻辑（需要计算输出对输入的导数）最适合用自定义训练循环实现，自由度更高也不会触发静态图占位符冲突，修改后的可运行代码如下：

from numpy import loadtxt
from keras.models import Sequential
from keras.layers import Dense
import tensorflow as tf #tf.__version__ = '2.6.0'
# 加载数据集
dataset = loadtxt('pima-indians-diabetes.csv', delimiter=',')
# 拆分输入输出
X = dataset[:,0:8] 
y = dataset[:,8]
X = tf.convert_to_tensor(X, dtype=tf.float32)
y = tf.convert_to_tensor(y, dtype=tf.float32)
# 构建模型
model = Sequential()
model.add(Dense(12, input_dim=8, activation='relu'))
model.add(Dense(12, activation='relu'))
model.add(Dense(12, activation='relu'))
model.add(Dense(1, activation='sigmoid'))
# 定义优化器和指标
optimizer = tf.keras.optimizers.Adam()
loss_metric = tf.keras.metrics.Mean(name='loss')
acc_metric = tf.keras.metrics.BinaryAccuracy(name='accuracy')
# 自定义训练步
@tf.function
def train_step(x, y_true):
    with tf.GradientTape(persistent=True) as t:
        t.watch(x)
        y_pred = model(x, training=True)
        # 计算PDE残差损失
        DyDX = t.gradient(y_pred, x)
        dy_t = DyDX[:, 5:6]
        loss_PDE = tf.reduce_mean(tf.square(dy_t))
        # 需添加数据损失可在此处写逻辑
        # loss_data = tf.reduce_mean(tf.square(y_true - y_pred))
        total_loss = loss_PDE
    # 计算梯度更新参数
    grads = t.gradient(total_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    # 更新训练指标
    loss_metric.update_state(total_loss)
    acc_metric.update_state(y_true, y_pred)
    return {
        'loss': loss_metric.result(),
        'accuracy': acc_metric.result()
    }
# 训练循环
EPOCHS = 15
for epoch in range(EPOCHS):
    # 重置每轮指标
    loss_metric.reset_states()
    acc_metric.reset_states()
    
    # 执行训练步，大数据集可自行修改为按batch迭代
    metrics = train_step(X, y)
    
    print(f'Epoch {epoch+1}/{EPOCHS}, loss: {metrics["loss"].numpy():.4f}, accuracy: {metrics["accuracy"].numpy():.4f}')

替代方案：使用add_loss适配fit模式 #

如果你不想完全改写训练逻辑，也可以通过add_loss接口将输入张量传入损失计算逻辑，修改如下：

from numpy import loadtxt
from keras.layers import Dense, Input
import tensorflow as tf

dataset = loadtxt('pima-indians-diabetes.csv', delimiter=',')
X = dataset[:,0:8] 
y = dataset[:,8]
X = tf.convert_to_tensor(X, dtype=tf.float32)
y = tf.convert_to_tensor(y, dtype=tf.float32)
# 显式定义输入层
inputs = Input(shape=(8,))
x = Dense(12, activation='relu')(inputs)
x = Dense(12, activation='relu')(x)
x = Dense(12, activation='relu')(x)
outputs = Dense(1, activation='sigmoid')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
# 计算PDE损失并添加到模型
with tf.GradientTape() as t:
    t.watch(inputs)
    pred = model(inputs)
DyDX = t.gradient(pred, inputs)
dy_t = DyDX[:,5:6]
pde_loss = tf.reduce_mean(tf.square(dy_t))
model.add_loss(pde_loss)
# 编译训练，此处loss填None即可，损失已通过add_loss添加
model.compile(optimizer='adam', metrics=['accuracy'])
model.fit(X, y, epochs=15)

内容的提问来源于stack exchange，提问作者Milad S