在 tensorflow 中使用 batch_jacobian 计算二阶导数在训练期间真的很慢

Question

我正在尝试计算神经网络输出相对于其输入的 Hessian。为了给你一个想法，这是我要计算的矩阵：

我是运行 Tensorflow 2.5.0，我计算 M 矩阵的代码如下所示：

def get_Mass_Matrix(self, q, dq):
    nDof = dq.shape[1]
    with tf.GradientTape(persistent = True) as t2:
        t2.watch(dq)
        with tf.GradientTape(persistent = True) as t1:
            t1.watch(dq)
            T = self.kinetic(q, dq)
            
        g = t1.gradient(T, dq)
    h = t2.batch_jacobian(g, dq)
        
    return h

函数self.kinetic()调用多层感知器。当我像这样计算 M 时，我得到了正确的答案，但我的神经网络训练显着减慢，即使在 GPU 上运行也是如此。

我想知道是否有更有效的方法来执行相同的计算而不会导致太多开销？谢谢。

作为参考，我正在使用子类化方法来构建模型（它继承自 tf.keras.Model）。

编辑：

添加有关 self.kinetic 函数的更多详细信息：

def kinetic(self, q, qdot):
    nDof = q.shape[1]
    qdq = tf.concat([tf.reshape(q, ((-1, nDof))),
                      tf.reshape(qdot, ((-1, nDof)))], axis = -1)
    
    return self.T_layers(qdq)

T_layers 定义为：

    self.T_layers = L(nlayers = 4, n = 8, input_dim = (latent_dim, 1), nlact = 'swish', oact = 'linear')

哪个正在调用：

class L(tf.keras.layers.Layer):

    def __init__(self, nlayers, n, nlact, input_dim, oact = 'linear'):

        super(L, self).__init__()

        self.layers = nlayers
        self.dense_in = tf.keras.layers.Dense(n, activation = nlact, input_shape = input_dim)
        self.dense_lays = []

        for lay in range(nlayers):
            self.dense_lays.append(tf.keras.layers.Dense(n, activation = nlact, kernel_regularizer = 'l1'))

        self.dense_out = tf.keras.layers.Dense(1, activation = oact, use_bias = False)

    def call(self, inputs):
        x = self.dense_in(inputs)
        for lay in range(self.layers):
            x = self.dense_lays[lay](x)

        return self.dense_out(x)

我怀疑部分问题可能是我没有“构建”图层？任何建议表示赞赏！

Answer 1

为了从 tensorflow 获得合理的性能，尤其是在计算梯度时，您必须用 @tf.function 修饰您的 get_Mass_Matrix 以确保它以图形模式运行。为此，函数内的所有内容都必须与图形模式兼容。

在class L的call函数中，最好直接遍历列表而不是索引，即：

class L(tf.keras.layers.Layer):
    ...
    def call(self, inputs):
        x = self.dense_in(inputs)
        for l in self.dense_lays:
            x = l(x)

        return self.dense_out(x)

然后，你可以装饰你的get_Mass_Matrix。

@tf.function
def get_Mass_Matrix(self, q, dq):
    with tf.GradientTape() as t2:
        t2.watch(dq)
        with tf.GradientTape() as t1:
            t1.watch(dq)
            T = self.kinetic(q, dq)    
        g = t1.gradient(T, dq)
    return t2.batch_jacobian(g, dq)

备注：传入get_Mass_Matrix的q和dq必须是常量[=30的张量 =] shape(constant between calls)，否则，每次有新的形状时它都会回溯并减慢速度。

在 tensorflow 中使用 batch_jacobian 计算二阶导数在训练期间真的很慢

Computation of second derivatives with batch_jacobian in tensorflow is really slow during training

python

tensorflow

tensorflow2.x