TensorFlow 二阶导数计算在 Quantum PINN 求解器中始终返回零

Question

我正在尝试实现一个简单的物理信息神经网络（PINN），用于在给定任意“标量势”的情况下求解“与时间无关的薛定谔方程”。我认为这些细节并不重要。我正在使用 TensorFlow 2.x 并遇到一个问题，即二阶导数计算始终返回 None，无论使用什么势函数。这是我的代码的简化版本，它显示了该问题：

import tensorflow as tf

class QuantumPINNSolver:
    def __init__(self, potential_func):
        self.potential_func = potential_func
        self.model = self.build_model()

    def build_model(self):
        return tf.keras.Sequential([
            tf.keras.layers.Dense(50, activation='relu', input_shape=(1,), dtype=tf.float64),
            tf.keras.layers.Dense(50, activation='relu', dtype=tf.float64),
            tf.keras.layers.Dense(50, activation='relu', dtype=tf.float64),
            tf.keras.layers.Dense(1, activation=None, dtype=tf.float64)
        ])

    @tf.function
    def schrodinger_loss(self, x, E):
        with tf.GradientTape(persistent=True) as tape2:
            with tf.GradientTape(persistent=True) as tape1:
                tape1.watch(x)
                psi = self.model(x)
            dpsi_dx = tape1.gradient(psi, x)
        d2psi_dx2 = tape2.gradient(dpsi_dx, x)
        
        print("psi:", psi)
        print("dpsi_dx:", dpsi_dx)
        print("d2psi_dx2:", d2psi_dx2)  # This is always zero
        
        V_x = self.potential_func(x)
        schrodinger_eq = -d2psi_dx2 + (V_x - E) * psi
        return tf.reduce_mean(tf.square(schrodinger_eq))

# Usage
def harmonic_oscillator_potential(x):
    return 0.5 * tf.square(x)

solver = QuantumPINNSolver(potential_func=harmonic_oscillator_potential)
x = tf.random.uniform((100, 1), minval=-5, maxval=5, dtype=tf.float64)
E = tf.constant(0.5, dtype=tf.float64)
loss = solver.schrodinger_loss(x, E)

问题是

d2psi_dx2

 始终是 

None

，即使

psi

和

dpsi_dx

具有非零值。我试过：

使用不同的势函数

更改模型架构

1. 使用不同的激活函数
2. 调整输入范围和
dtype

4. ...但二阶导数仍然是

   None
。可能是什么原因造成的？如何修复它以正确计算二阶导数？我对

TensorFlow

不太熟悉（而且我仍在努力理解

GradientTape

），我希望这是一些我通过阅读文档尚未发现的简单问题。 想通了！

Answer 1

GradientTape

旨在计算标量输出相对于所观察变量的梯度。然而，显然当您尝试计算高阶导数（如二阶导数）时，TensorFlow 需要显式处理。

在原始代码中，第二个梯度 (

d2psi_dx2

) 返回 None，因为 TensorFlow 的梯度跟踪不会自动通过嵌套磁带传播梯度以获得高阶导数。

要解决此问题并正确计算二阶导数，您需要使用单独的 GradientTape 上下文显式管理梯度计算：

使用多个 GradientTape 上下文：

使用一个
1. GradientTape

   计算 psi，使用另一个嵌套 GradientTape 计算 dpsi_dx。

   持久磁带：
确保两个
2. GradientTape

   上下文都声明为持久 (persistent=True)，以便您可以在其中多次计算梯度。

   显式导数计算：
计算一阶导数 (
3. dpsi_dx

   ) 后，使用外部 GradientTape 上下文计算相对于输入

   d2psi_dx2

   的二阶导数 (

   x

   )。

   def schrodinger_loss(self, x, E):
           with tf.GradientTape(persistent=True) as tape:
               tape.watch(x)
               with tf.GradientTape(persistent=True) as tape1:
                   tape1.watch(x)
                   psi = self.model(x)
               dpsi_dx = tape1.gradient(psi, x)
           
           d2psi_dx2 = tape.gradient(dpsi_dx, x)
           
           V_x = self.potential_func(x)
           schrodinger_eq = -d2psi_dx2 + (V_x - E) * psi
           return tf.reduce_mean(tf.square(schrodinger_eq))