TensorFlow Keras自定义损失函数报错及散度计算问题求助

错误核心原因 #

你遇到的问题根源有两点：

• 自定义损失函数中直接调用了非TensorFlow原生操作（如y_pred.numpy()、scipy的Triangulation和LinearTriInterpolator），这些操作无法被TensorFlow的符号图追踪，即便开启run_eagerly=True，也会因混合符号张量（KerasTensor）与Eager操作触发报错。
• 你将模型的Input符号张量x传入损失函数custom_loss(x)，导致Keras构建符号图时，把损失函数与模型输入张量绑定，触发不支持的符号操作调用。

解决方案步骤 #

1. 用tf.py_function包装非TensorFlow操作

把依赖scipy的散度计算逻辑封装为Python函数，再通过tf.py_function转为TensorFlow可识别的操作，确保计算能被Eager模式正确处理（若需要梯度传递，需额外自定义梯度逻辑）。

2. 避免在损失函数中使用模型的Input符号张量

损失函数仅依赖y_true和y_pred，外部数据（如X、Y、网格信息）通过闭包参数传入，而非模型的Input张量。若需使用模型输入数据，可将其作为模型的额外输出，训练时同步传入。

3. 修正张量类型与转换问题

移除不必要的tf.constant转换，统一张量数据类型（如均使用tf.float32），避免类型不匹配错误。

修改后的代码示例 #

import tensorflow as tf
import numpy as np
from scipy.interpolate import LinearTriInterpolator
from scipy.spatial import Delaunay

# 提前构建三角剖分，避免在损失函数中重复初始化
tri = Delaunay(np.vstack((X, Y)).T)

def compute_divergence(f_np, X, Y, tri):
    # 处理numpy数组的散度计算逻辑
    Fx = LinearTriInterpolator(tri, f_np[:, 0, 0] + f_np[:, 0, 2])
    Fy = LinearTriInterpolator(tri, f_np[:, 0, 1] + f_np[:, 0, 2])
    gradx = Fx.gradient(X, Y)
    grady = Fy.gradient(X, Y)
    div = gradx[0] + grady[0] + gradx[1] + grady[1]
    return div.flatten()

def custom_loss(X, Y, tri):
    def loss(y_true, y_pred):
        # 用tf.py_function包装散度计算，适配TensorFlow Eager模式
        div = tf.py_function(
            func=lambda f: compute_divergence(f.numpy(), X, Y, tri),
            inp=[y_pred],
            Tout=tf.float32
        )
        L_div = tf.reduce_mean(tf.math.square(div))

        # 修正积分误差计算：若需使用输入strains，从模型额外输出中获取
        # 这里假设模型输出包含strains，对应训练时传入的第二个y参数
        strains = y_true[1]
        Wi = tf.reduce_sum(tf.linalg.matmul(tf.expand_dims(strains, 1), tf.transpose(y_pred[0], perm=[0, 2, 1])), axis=0)
        Wext = tf.reduce_mean(y_true[0], axis=0)
        L_en = tf.abs(Wi - Wext)
        L_en = tf.cast(L_en, dtype=tf.float32)

        # 总损失：统一数据类型
        total_loss = 0.095 * L_div + 0.05 * L_en
        return total_loss
    return loss

# 修改模型：将输入strains作为额外输出，方便损失函数调用
x = tf.keras.Input(shape=(3,), name='strains')
x_t = tf.keras.layers.Reshape((1, 3))(x)
hidden = x_t
for i in range(n_layers):
    hidden = tf.keras.layers.Dense(units=n_neurons, activation=act_fun)(hidden)
f = tf.keras.layers.Dense(units=3, activation='relu', name='out')(hidden)
# 新增输出：返回输入的strains
model = tf.keras.Model(inputs=x, outputs=[f, x])

# 编译模型：传入外部网格数据，无需传入模型Input张量
model.compile(
    loss=custom_loss(X, Y, tri),
    optimizer=tf.keras.optimizers.Adam(),
    metrics=[],
    run_eagerly=True
)

# 训练：对应模型的两个输出，传入[y_train, x_train]
early_stop = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=5)
model.fit(
    x_train, 
    [y_train, x_train],
    validation_data=(x_test, [y_test, x_test]),
    batch_size=2500,
    epochs=100,
    shuffle=False, 
    callbacks=[early_stop],
    verbose=1
)

额外注意事项 #

• 若compute_divergence需要梯度传递，需用tf.custom_gradient装饰该函数，实现自定义梯度逻辑（tf.py_function默认不追踪梯度）。
• 提前初始化三角剖分可避免重复计算，提升训练效率。
• 保持所有张量数据类型一致，避免float32与float64混合导致的类型错误。

内容的提问来源于stack exchange，提问作者Alberto Ciampaglia