Разреженная матрица треугольного решения в Tensorflow?

Question

Была ли реализация для решения Ax = b с разреженной треугольной матрицей в Tensorflow? (соответствует tf.matrix_triangular_solve ())

AFAIK, если мы имеем A, например, как нижнюю треугольную матрицу с представлением разреженной матрицы, нам нужно преобразовать ее в плотную матрицу с помощью tf.sparse_to_dense().

Но если A имеет очень большое измерение, например, 16416x16416, и очень разреженные записи, например, 0,018% (около 45216 ненулей), это заняло бы большой объем памяти.

Я думаю, было бы очень полезно, если бы мы могли воспользоваться преимуществами разреженного матричного решателя, такого как матрица с полосчатой ​​структурой, в Tensorflow.

Извините, если мой вопрос не актуален. Например, если бы были какие-то решения для этого, я был бы благодарен за любую помощь.

Благодарю.

Answer 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

По крайней мере, для меня решение очень быстрое. Дайте мне знать, если вы обнаружите ошибку (например, в вычислениях градиента)

Answer 2

В TF очень мало поддержки разреженных тензоров. Таким образом, ваш единственный подход в настоящее время (как вы определили) это tf.sparse_to_dense()