Проблема с памятью: пакетное построение матрицы нижнего tri angular (векторизованным способом)

Question

Я пытаюсь достичь следующего:

Дан массив, который выглядит так:

[1 0 0 0 0 0 0 0 0]
[1 1 0 0 0 0 0 0 0]
[1 1 1 0 0 0 0 0 0]
[1 1 1 1 0 0 0 0 0]
[1 1 1 1 1 0 0 0 0]
[1 1 1 1 1 1 0 0 0]
[1 1 1 1 1 1 1 0 0]
[1 1 1 1 1 1 1 1 0]
[1 1 1 1 1 1 1 1 1]

Имейте в виду, что это небольшой массив формы [9 на 9] и может быть инициализирован одним go. Массив, который я хочу построить, имеет форму [90000 на 90000], следовательно, его нельзя инициализировать в одном go.

После преобразования этого массива [9 на 9] в партии он выглядит так:

    [[1 0 0 0 0 0 0 0 0]
    [1 1 0 0 0 0 0 0 0]
    [1 1 1 0 0 0 0 0 0]],

    [[1 1 1 1 0 0 0 0 0]
    [1 1 1 1 1 0 0 0 0]
    [1 1 1 1 1 1 0 0 0]],

    [[1 1 1 1 1 1 1 0 0]
    [1 1 1 1 1 1 1 1 0]
    [1 1 1 1 1 1 1 1 1]]

Я могу использовать эти отдельные фрагменты для выполнения операций.

Как я могу инициализировать массив [90000 на 90000] партиями, сохраняя размещение единиц в массиве?

Я пробовал tf.linalg.LinearOperatorLowerTrin angular () оператор. Вот как выглядит результат:

>>> arr = tf.linalg.LinearOperatorLowerTriangular(tf.ones((3,3,9)))
>>> arr.to_dense()
<tf.Tensor: id=28, shape=(3, 3, 9), dtype=float32, numpy=
array([[[1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0.]],

       [[1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0.]],

       [[1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0.]]], dtype=float32)>

На выходе вы можете видеть, что единицы в каждом пакете начинаются с первого индекса.

Кроме того, когда я инициализирую "arr" :

>>> arr = tf.linalg.LinearOperatorLowerTriangular(tf.ones((3,3,9)))

Я использовал tf.ones(), в моем случае я даже не могу использовать это для инициализации массива 90000 на 90000.

Пожалуйста, предложите какой-либо эффективный способ создания / инициализации такого матрица?

Редактировать:

Что именно я ищу? a logi c относительно того, как я могу сделать блоки / субматрицы большой матрицы 90000 на 90000. Например, я могу создать 10 субматриц размером 9000 на 90000. Однако как я могу инициализировать эти 10 субматриц, чтобы выглядят как единственная нижняя три angular матрица. Я буду использовать эти подматрицы для выполнения векторно-матричного умножения. Как только я смогу найти способ присвоить значения этим подматрицам (таким образом, чтобы при сложении они выглядели как одна нижняя три angular матрица единиц), то после выполнения векторно-матричного умножения я могу добавить все 10 результирующих матриц. в один.

Answer 1

Я не совсем уверен, что это то, что вам нужно, но вот способ вычислить что-то подобное по кускам в al oop в TensorFlow (здесь 2.x, но должно быть аналогично для 1.x) :

import tensorflow as tf

@tf.function
def mult_tril_chunks(vector, num_chunks=None, chunk_size=None):
    # Function accepts either number of chunks or size of each chunk
    # Last chunk may have different size
    if num_chunks is None == chunk_size is None:
        raise ValueError('one and only one of num_chunks and chunk_size must be given.')
    vector = tf.convert_to_tensor(vector)
    # Compute number of chunks or chunk size depending on given parameters
    s = tf.shape(vector)[0]
    if chunk_size is not None:
        chunk_size = tf.dtypes.cast(chunk_size, s.dtype)
        num_chunks = (s // chunk_size) + tf.dtypes.cast(s % chunk_size != 0, s.dtype)
    else:
        n = tf.dtypes.cast(num_chunks, s.dtype)
        chunk_size = s // n + tf.dtypes.cast(s % n != 0, s.dtype)
    # If you know all chunks will always have the same size
    # you can instead pass element_shape=[chunk_size]
    ta = tf.TensorArray(vector.dtype, size=num_chunks, element_shape=[None],
                        infer_shape=False)
    # Do computation in a loop
    _, ta = tf.while_loop(
        lambda i, ta: i < num_chunks,
        lambda i, ta: _mult_tril_chunks_loop(i, ta, vector, s, chunk_size),
        [0, ta]
    )
    # Return concatenated result
    return ta.concat()

def _mult_tril_chunks_loop(i, ta, vector, s, chunk_size):
    # Chunk bounds
    start = i * chunk_size
    end = tf.math.minimum(start + chunk_size, s)
    # Make slice of lower triangular matrix
    r = tf.range(s)
    tril_slice = r <= tf.expand_dims(tf.range(start, end), axis=1)
    tril_slice = tf.dtypes.cast(tril_slice, vector.dtype)
    # Compute product
    mult_slice = tf.linalg.matvec(tril_slice, vector)
    return i + 1, ta.write(i, mult_slice)

# Test
vector = tf.range(10)
# Result with normal computation
tril = tf.linalg.band_part(tf.ones((10, 10), vector.dtype), -1, 0)
res = tf.linalg.matvec(tril, vector)
print(res.numpy())
# [ 0  1  3  6 10 15 21 28 36 45]
# Giving number of chunks
res2 = mult_tril_chunks(vector, num_chunks=3)
print(res2.numpy())
# [ 0  1  3  6 10 15 21 28 36 45]
# Giving chunk size
res3 = mult_tril_chunks(vector, chunk_size=4)
print(res3.numpy())
# [ 0  1  3  6 10 15 21 28 36 45]

Поскольку вас беспокоит использование памяти, возможно, стоит проверить некоторые из более сложных параметров для tf.while_loop, например, количество параллельных итераций для запуска .