TensorFlow自动微分中线性网络二阶导数为None的问题及诉求

问题原因 #

全采用线性激活的神经网络，本质是多层线性变换的组合，最终输出可简化为输入的线性函数（形如 y = a*x + b）。这类函数的一阶导数是常数，二阶导数则恒为0。

TensorFlow的GradientTape在计算梯度时会做优化：如果检测到导数与被求导变量完全无关（即梯度恒为0），会直接返回None，而非生成全0张量，以此减少不必要的计算开销和内存占用。这就是y_xx_lin返回None的核心原因；而使用tanh激活的网络因为存在非线性变换，二阶导数与输入相关，所以能得到具体数值。

解决办法：强制返回全0张量 #

要让TensorFlow返回实际的0值而非None，只需在计算二阶导数后添加判断，若结果为None则手动生成与输入形状一致的全0张量。修改后的get_first_second_derivative函数如下：

def get_first_second_derivative(X_train,y_train,model):
    with tf.GradientTape(persistent=True) as tape_second:
        tape_second.watch(X_train)
        
        with tf.GradientTape(persistent=True) as tape_first:
            tape_first.watch(X_train)
            output = model(X_train)
    
        y_x  = tape_first.gradient(output,X_train)

    y_xx = tape_second.gradient(y_x,X_train)
    # 若二阶导数为None，生成与输入同形状的全0张量
    if y_xx is None:
        y_xx = tf.zeros_like(X_train)
    
    return y_x,y_xx

修改后运行代码，y_xx_lin就会变成和X_train形状一致的全0张量，而非None。

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