TensorFlow Batch Hessian

Question

我正在构建一个必须逼近某个多元函数的神经网络，例如 f(x)。损失函数定义为网络的二阶导数与函数f的接近程度。为此，我必须计算 f(x) 的 Hessian 矩阵。我编写了一个自定义 TensorFlow 模型，看起来像这样

class ApproximateModel(tf.keras.Model):
    @tf.function
    def f_true_hessian(x: tf.Tensor) -> tf.Tensor:
        # Some function that should return the actual Hessian
        return x
    
    def train_step(self, data):
        # tf.get_shape(x) -> (batch_size, dimension_x)
        x = data[0]
        
        # Calculate loss
        with tf.GradientTape() as second_tape:
            with tf.GradientTape() as first_tape:
                first_tape.watch(x)
                second_tape.watch(x)
                
                f = self(x, training=True)
            
            f_x = first_tape.gradient(f, x)
            second_tape.watch(f_x)
        
        f_jacobian = second_tape.jacobian(f_x, x)
        # tf.get_shape(f) -> (batch_size, dimension_x, batch_size, dimension_x)
        
        # I want to get (batch_size, dimension_x, dimension_x) somehow..
        loss = tf.math.reduce_mean(tf.math.square(tf.reduce_sum(f_jacobian, axis=[1, 2]) - self.f_true_hessian(x))))
        return loss

对于感兴趣的读者，这种类型的网络的应用是近似偏微分方程，如这里。

如果我没有批量大小，上面的代码可以很好地工作。如果我有一批 x 的样本，我不知道如何获得 Hessian 矩阵。如何获得所需的输出，其中仅计算 dimension_x 的 Hessian 矩阵并省略 batch_size？

Answer 1

我刚刚发现自己处于类似的情况。我不知道这是否是最好的解决方案，但我所做的是堆叠渐变

f_jacobian = tf.stack([second_tape.gradient(f_x[:,i:i+1], x) for i in range(dimension_x)], 1)