在tensorflow中用batch_jacobian计算二阶导数在训练过程中真的很慢

如何解决在tensorflow中用batch_jacobian计算二阶导数在训练过程中真的很慢

我正在尝试计算神经网络输出相对于其输入的 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,qdot):
    nDof = q.shape[1]
    qdq = tf.concat([tf.reshape(q,((-1,nDof))),tf.reshape(qdot,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,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)

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

解决方法

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

在call的{​​{1}}函数中，最好直接迭代列表而不是索引它，即：

class L

然后，您可以装饰您的 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) 的 q 和 dq 必须是 tensor 的 constant 形状（调用之间的常量）,否则每次出现新形状时都会回溯并减速。