当行列式 overflows/underflows 时计算 TensorFlow 中行列式的对数

Calculate log of determinant in TensorFlow when determinant overflows/underflows

我的成本函数涉及 log(det(A)) 的计算(假设 det(A) 为正,因此对数有意义,但 A 不是 Hermitian,因此 Cholesky 分解不适用于此处)。当 det(A) 非常 large/small 时,直接调用 det(A) 将 overflow/underflow。为了避免这种情况,可以使用

的数学事实

log(det(A)) = Tr(log(A)),

其中后者可以使用 LU 分解进行评估(比 eigenvalue/SVD 更有效)。该算法已经在numpy中实现为numpy.linalg.slogdet,所以问题是如何从TensorFlow中调用numpy。


这是我试过的

import numpy as np
import tensorflow as tf
from tensorflow.python.framework import function

def logdet_np(a):
    _, l = np.linalg.slogdet(a)
    return l

def logdet1(a):
    return tf.py_func(logdet_np, [a], tf.float64)

@function.Defun(tf.float64, func_name='LogDet')
def logdet2(a):
    return tf.py_func(logdet_np, [a], tf.float64)

with tf.Session() as sess:
    a = tf.constant(np.eye(500)*10.)
    #print(sess.run(logdet1(a)))
    print(sess.run(logdet2(a)))

我先定义一个python函数来传递numpy的结果。然后我使用 tf.py_func 定义了两个 logdet 函数。第二个函数由 function.Defun 装饰,用于稍后定义 TensorFlow 函数及其梯度。当我测试它们时,我发现第一个函数 logdet1 有效并给出了正确的结果。但是第二个函数logdet2returns一个KeyError.

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/Library/Frameworks/Python.framework/Versions/3.5/lib/python3.5/site-
packages/tensorflow/python/ops/script_ops.py in __call__(self, token, args)
     77   def __call__(self, token, args):
     78     """Calls the registered function for `token` with args."""
---> 79     func = self._funcs[token]
     80     if func is None:
     81       raise ValueError("callback %s is not found" % token)

KeyError: 'pyfunc_0'

我的问题是 Defun 装饰器有什么问题?为什么和py_func冲突?如何在 TensorFlor 中正确包装 numpy 函数?


logdet 定义梯度的剩余部分与问题 有关。根据该问题的解决方案,尝试编写

@function.Defun(tf.float64, tf.float64, func_name='LogDet_Gradient')
def logdet_grad(a, grad):
    a_adj_inv = tf.matrix_inverse(a, adjoint=True)
    out_shape = tf.concat([tf.shape(a)[:-2], [1, 1]], axis=0)
    return tf.reshape(grad, out_shape) * a_adj_inv
@function.Defun(tf.float64, func_name='LogDet', grad_func=logdet_grad)
def logdet(a):
    return tf.py_func(logdet_np, [a], tf.float64, stateful=False, name='LogDet')

如果能解决Defunpy_func之间的冲突,上面的代码就可以工作了,这是我上面提出的关键问题。

如果你的问题是溢出,你可以用简单的数学来解决它。

所以你只需要获取特征值,记录它们并求和。

可以用SVD分解A:

A = U S V'

由于一个乘积的行列式是行列式的乘积,UV'的行列式为1或-1,而S的行列式为非-否定,

abs(det(A)) = det(S)

因此,(正)行列式的对数可以计算为

tf.reduce_sum(tf.log(svd(A, compute_uv=False)))


从TF1.1开始,tf.svd缺少渐变(未来的版本可能会有),所以我建议采用kofd的代码实现:

https://github.com/tensorflow/tensorflow/issues/6503

在@MaxB 的帮助下,我在这里 post 定义函数的代码 logdet 用于 log(abs(det(A))) 及其梯度。

  • logdet 调用 numpy 函数 numpy.linalg.slogdet 使用 log(det(A))=Tr(log(A)) 的思想计算行列式的对数,这对行列式的 overflow/underflow 是稳健的。它基于 LU 分解,与基于 eigenvalue/SVD 的方法相比效率更高。

  • numpy 函数 slogdet returns 包含行列式符号和 log(abs(det(A))) 的元组。该符号将被忽略,因为它不会对优化中的梯度信号产生影响。

  • logdet的梯度是通过矩阵求逆计算的,根据grad log(det(A)) = inv(A)^T。它基于 _MatrixDeterminantGrad 上的 TensorFlow 代码,稍作修改。


import numpy as np
import tensorflow as tf
# from https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342
# Define custom py_func which takes also a grad op as argument:
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):
    # Need to generate a unique name to avoid duplicates:
    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))
    tf.RegisterGradient(rnd_name)(grad)  # see _MySquareGrad for grad example
    g = tf.get_default_graph()
    with g.gradient_override_map({"PyFunc": rnd_name}):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)
# from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_grad.py
# Gradient for logdet
def logdet_grad(op, grad):
    a = op.inputs[0]
    a_adj_inv = tf.matrix_inverse(a, adjoint=True)
    out_shape = tf.concat([tf.shape(a)[:-2], [1, 1]], axis=0)
    return tf.reshape(grad, out_shape) * a_adj_inv
# define logdet by calling numpy.linalg.slogdet
def logdet(a, name = None):
    with tf.name_scope(name, 'LogDet', [a]) as name:
        res = py_func(lambda a: np.linalg.slogdet(a)[1], 
                      [a], 
                      tf.float64, 
                      name=name, 
                      grad=logdet_grad) # set the gradient
        return res

可以测试 logdet 适用于非常 large/small 的行列式并且它的梯度也是正确的。

i = tf.constant(np.eye(500))
x = tf.Variable(np.array([10.]))
y = logdet(x*i)
dy = tf.gradients(y, [x])
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run([y, dy]))

结果:[1151.2925464970251, [array([ 50.])]]