Tensorflow中的稀疏矩阵三角形求解?

时间:2017-07-12 08:26:54

标签: memory tensorflow sparse-matrix linear-equation

是否有在Tensorflow中使用稀疏三角矩阵求解Ax = b的实现? (对应于tf.matrix_triangular_solve())

AFAIK,如果我们有A,例如,作为具有稀疏矩阵表示的下三角矩阵,我们需要使用tf.sparse_to_dense()将其转换为密集矩阵。

但是,如果A具有非常大的维度,例如16416x16416,并且非常稀疏的条目,例如0.018%(大约45216个非零),则需要大量的内存。

我认为如果我们可以在Tensorflow中利用稀疏矩阵求解器,例如带有带状结构的矩阵,将会非常有用。

抱歉,如果我的问题无关紧要。 例如,如果有任何解决方案,我将不胜感激任何帮助。

感谢。

2 个答案:

答案 0 :(得分:1)

我有同样的问题,我为它创建了一个自定义操作。只要您不想计算渐变到A和A仍然是固定的,那么这段代码应该有所帮助:

import tensorflow as tf
import numpy as np
from scipy.sparse import linalg as sla
import scipy

lu = sla.splu(A)

# Define custom py_func which takes also a grad op as argument:
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):

    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

    tf.RegisterGradient(rnd_name)(grad)
    g = tf.get_default_graph()
    with g.gradient_override_map({"PyFunc": rnd_name}):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)


def sparse_solve(x, lu, dtype=tf.float64, name=None):

    with tf.name_scope(name, 'SparseSolve', [x]) as name:
        solve_x = py_func(lu.solve,
                        [x],
                        [dtype],
                        name=name,
                        grad=_SparseSolveGrad(dtype, lu))
        return solve_x[0]

class _SparseSolveGrad:
    def __init__(self, dtype, lu):
        self.dtype = dtype
        self.lu = lu

    def __call__(self, op, grad):
        x = op.inputs[0]
        y = tf.conj(tf.py_func(self.lu.solve, [tf.conj(grad)], self.dtype))
        return y

解决方案基于我在https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342

找到的代码

至少对我而言,解决方案非常快。如果您发现错误(例如在渐变计算中),请告诉我

答案 1 :(得分:0)

TF中有very little support for sparse tensors。因此,您目前唯一的方法(正如您所确定的)是tf.sparse_to_dense()