使用 TensorFlow v2 计算经过训练的神经网络的偏导数

Question

我想计算由深度神经网络近似的函数的一阶和二阶导数。我尝试了一个简单的例子，我们可以清楚地看到正确的答案是什么： 假设神经网络要逼近的函数为：y=x^2 然后，我按如下方式构建网络：

import numpy as np
import math
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import losses

m_test=np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14])
n_test= m_test**2

def model_gen(Input_shape):
    X_input = Input(shape=Input_shape)
    X = Dense(units=128, activation='relu')(X_input)
    X = Dense(units=64, activation='relu')(X)
    X = Dense(units=32, activation='relu')(X)
    X = Dense(units=1)(X)
    model = Model(inputs=X_input, outputs=X)
    return model

model = model_gen(Input_shape=(1,))
opt = Adam(lr=0.01, beta_1=0.9, beta_2=0.999, decay=0.001)
model.compile(loss=losses.mean_squared_error, optimizer=opt)
model.fit(m_test,n_test, epochs=1000)

训练模型后，我使用 GradientTape() 计算导数，如下所示：

x = np.array([1,2,3,4,5,6])
x = tf.convert_to_tensor(x, np.float32)
with tf.GradientTape() as tape:
    tape.watch(x)
    with tf.GradientTape() as tape2:
        tape2.watch(x)
        y = model(x)
    grad_y = tape2.gradient(y, x)
grad_y2 = tape.jacobian(grad_y,x)    
print(y)
print(grad_y)
print(grad_y2)

结果如下：

tf.Tensor(
[[ 1.1044241]
 [ 4.0293255]
 [ 8.592318 ]
 [16.362782 ]
 [24.958858 ]
 [36.002636 ]], shape=(6, 1), dtype=float32)
tf.Tensor([ 0.9645251  2.2506325  7.7348967  7.901049   9.065702  11.975435 ], shape=(6,), dtype=float32)
tf.Tensor(
[[0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0.]], shape=(6, 6), dtype=float32)

虽然一阶导数在某种程度上可以接受（存在错误），但二阶导数为 0，这绝对是错误的，因为我们知道函数 y=x^2 的二阶导数始终为 2。

如果有人可以帮助我，我将不胜感激。多谢。