TensorFlow自定义PDE求解损失函数及梯度无值问题排查

错误根源 #

出现"No gradients provided for any variable"的核心原因是：计算t置0后的预测值时，操作脱离了TensorFlow自动微分的追踪链路——比如用了numpy操作修改输入、未在GradientTape上下文内调用模型，或者使用了不可微分的张量修改方式，导致梯度无法回传到模型权重。

正确实现方案 #

1. 模型定义（示例MLP） #

先定义一个基础的PDE求解模型，输入包含空间变量x和时间变量t：

import tensorflow as tf
from tensorflow.keras import layers

class PDEModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense_layers = [
            layers.Dense(64, activation='tanh'),
            layers.Dense(64, activation='tanh'),
            layers.Dense(1)
        ]
    
    def call(self, inputs):
        x = inputs
        for layer in self.dense_layers:
            x = layer(x)
        return x

2. 自定义可微分损失函数 #

损失函数需满足：所有模型调用、输入变换都在可微分链路内，用TensorFlow原生操作生成t=0的输入，不能破坏梯度追踪：

def pde_loss(model, inputs):
    # inputs: 形状(batch_size, 2)，[:,0]是x，[:,1]是t
    # 1. 计算原输入的预测值
    y_pred = model(inputs)
    
    # 2. 生成t=0的输入（用tf.concat实现可微分的张量修改）
    inputs_t0 = tf.concat([
        inputs[:, :1],  # 保留x不变
        tf.zeros_like(inputs[:, 1:])  # t替换为0
    ], axis=1)
    # 3. 计算t=0时的预测值（必须在梯度追踪范围内调用模型）
    y_pred_t0 = model(inputs_t0)
    
    # 4. 定义损失项（以Burgers方程为例，可根据你的PDE调整）
    # 初始条件损失：u(x,0) = sin(x)
    initial_loss = tf.reduce_mean(tf.square(y_pred_t0 - tf.sin(inputs[:, :1])))
    
    # PDE残差损失：u_t + u*u_x = 0
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(inputs)
        u = model(inputs)
        # 计算偏导数
        u_grad = tape.gradient(u, inputs)
        u_t = u_grad[:, 1:]  # 对t的偏导
        u_x = u_grad[:, :1]  # 对x的偏导
    pde_residual = u_t + u * u_x
    pde_loss = tf.reduce_mean(tf.square(pde_residual))
    
    # 总损失
    return initial_loss + pde_loss

3. 手动训练循环（关键） #

由于自定义损失需要传入模型本身，不能用Keras默认的compile()+fit()，必须手动用GradientTape追踪所有梯度相关操作：

# 初始化模型和优化器
model = PDEModel()
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

# 生成训练数据（示例：x∈[-π, π], t∈[0,1]）
batch_size = 64
x = tf.random.uniform((batch_size, 1), minval=-3.14159, maxval=3.14159)
t = tf.random.uniform((batch_size, 1), minval=0., maxval=1.)
train_inputs = tf.concat([x, t], axis=1)

# 训练步骤（用@tf.function加速）
@tf.function
def train_step(inputs):
    with tf.GradientTape() as tape:
        # 所有模型调用、损失计算必须在tape上下文内
        loss = pde_loss(model, inputs)
    # 计算模型权重的梯度
    grads = tape.gradient(loss, model.trainable_variables)
    # 更新权重
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# 启动训练
for epoch in range(100):
    current_loss = train_step(train_inputs)
    if epoch % 10 == 0:
        print(f"Epoch {epoch:3d} | Loss: {current_loss.numpy():.6f}")

常见错误规避 #

• 禁止用numpy修改输入后转回张量：这会断开梯度链，必须全程用TensorFlow操作（如tf.concat、tf.where）处理输入。
• 禁止在GradientTape外调用模型计算t=0的预测值：只有在tape上下文内的模型调用才会被追踪梯度。
• 不要用tf.assign直接修改张量元素：这类操作不可微分，必须创建新的张量来实现输入变换。

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