无监督神经网络求解定态Schrödinger方程代码报错修复求助

问题背景 #

复现论文《Neural-network-based multistate solver for a static Schrödinger equation》时，构建输入为x坐标、输出为N个波函数的神经网络，仅实现总能量损失项后运行报错，错误提示梯度计算时target为None。

错误原因 #

报错出现在kinetic_energy_term函数的二阶梯度计算中：

• 原代码中psi是模型的输出，它是神经网络权重的函数，而非输入x的可微分函数（x作为输入数据，默认不会被GradientTape追踪其与模型输出的梯度关系）。
• 第一次调用tape.gradient(psi, x)时，由于psi和x之间没有可追踪的依赖关系，返回None，导致第二次求二阶梯度时触发类型错误。

修复方案 #

核心是让模型输出psi成为x的可微分函数，需要在GradientTape内部完成模型对x的前向传播，确保梯度可以被正确追踪。以下是修改后的完整代码：

import tensorflow as tf
import numpy as np

# Constants
hbar = tf.constant(1.0, dtype=tf.float32)
m = tf.constant(1.0, dtype=tf.float32)
omega = tf.constant(1.0, dtype=tf.float32)

# Potential energy function
def potential_energy(x):
    return 0.5 * m * omega ** 2 * x ** 2

# 修改后的Kinetic energy term：接收模型和x，在tape内部计算psi
def kinetic_energy_term(model, x):
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(x)
        psi = model(x)  # 在tape内部计算模型输出，建立x和psi的梯度依赖
        dpsi_dx = tape.gradient(psi, x)
    d2psi_dx2 = tape.gradient(dpsi_dx, x)
    del tape
    return -0.5 * hbar ** 2 / m * d2psi_dx2

# 修改后的Expected energy function：计算所有态的能量
def expected_energy(model, x):
    psi = model(x)
    ke = kinetic_energy_term(model, x)
    pe = potential_energy(x)
    hamiltonian_psi = ke + pe * psi
    
    # 对每个态计算<psi|H|psi>/<psi|psi>
    numerators = tf.reduce_sum(psi * hamiltonian_psi, axis=0)
    denominators = tf.reduce_sum(psi * psi, axis=0)
    energies = numerators / denominators
    return energies

# Custom loss function：用TensorFlow操作代替Python循环，适配AutoGraph
def unsupervised_loss(model, x):
    energies = expected_energy(model, x)
    return tf.reduce_sum(energies)

# 自定义训练步骤，避免Keras默认fit的y_true冗余问题
@tf.function
def train_step(model, x, optimizer):
    with tf.GradientTape() as tape:
        loss = unsupervised_loss(model, x)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# Neural network model
N = 10  # Number of quantum states
model = tf.keras.Sequential([
    tf.keras.layers.Dense(50, activation='tanh', input_shape=(1,)),  # 替换sigmoid为tanh，避免边界饱和
    tf.keras.layers.Dense(50, activation='tanh'),
    tf.keras.layers.Dense(N, activation='linear')
])

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

# Training data
x_train = np.linspace(-5, 5, 1000).reshape(-1, 1)
x_train_tensor = tf.convert_to_tensor(x_train, dtype=tf.float32)

# 划分验证集
val_split = 0.2
val_size = int(len(x_train) * val_split)
x_val_tensor = x_train_tensor[-val_size:]
x_train_tensor = x_train_tensor[:-val_size]

# Train the model
epochs = 100
batch_size = 32
train_dataset = tf.data.Dataset.from_tensor_slices(x_train_tensor).shuffle(len(x_train_tensor)).batch(batch_size)
val_dataset = tf.data.Dataset.from_tensor_slices(x_val_tensor).batch(batch_size)

for epoch in range(epochs):
    # 训练
    train_loss = 0.0
    for batch_x in train_dataset:
        loss = train_step(model, batch_x, optimizer)
        train_loss += loss.numpy()
    train_loss /= len(train_dataset)
    
    # 验证
    val_loss = 0.0
    for batch_x in val_dataset:
        loss = unsupervised_loss(model, batch_x)
        val_loss += loss.numpy()
    val_loss /= len(val_dataset)
    
    print(f"Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

关键修改说明 #

1. 梯度追踪逻辑调整：将模型的前向传播psi = model(x)放入kinetic_energy_term的GradientTape内部，确保x和psi之间的梯度依赖被正确追踪。
2. 替换Python循环：用TensorFlow的向量化操作计算所有态的能量，避免AutoGraph处理Python循环时的兼容性问题。
3. 自定义训练循环：避开Kerasfit方法对y_true的强制要求，更适配无监督训练场景。
4. 激活函数优化：将sigmoid替换为tanh，避免在x边界处激活函数饱和导致梯度消失。

后续扩展建议 #

• 添加正交归一化损失项：计算不同态之间的内积，约束波函数的正交性和归一性，例如ortho_loss = tf.reduce_sum(tf.square(tf.matmul(psi, psi, transpose_a=True) - tf.eye(N)))。
• 添加L2正则化：在损失中加入模型权重的L2范数，防止过拟合，例如l2_loss = tf.add_n([tf.nn.l2_loss(var) for var in model.trainable_variables])，再乘以正则化系数后加入总损失。

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