Большой массив матричных точек при многопроцессорной обработке

Question

Я реализую особый случай EM-GMM.

X - это матрица данных формы [1000000, 900], и это объект nmap-mmap
Q - точная матрица формы [900, 900] и ndarray

Я также использую библиотеку Multiprocessing для одновременной обработки более 200 Q-матриц на 40 ядрах, используя одну и ту же матрицу данных (X).

Это работаетдля меньших размеров, таких как [1mil, 196], [1mil, 400],
, но когда я пытаюсь запустить [1mil, 900] в какой-то момент одного из процессов, возникает исключение:

OSError: [Errno 12] Невозможно выделить память

Думаю, проблема в том, что у меня есть 2 больших вычисления, которые, вероятно, выделяют большие матрицы.

Как часть E-шаг, который мне нужно рассчитать:
np.sum(X.dot(Q) * X, axis=1)

Как часть М-шага, мне нужно вычислить (W представляет собой вектор весов [1mil, 1]):
(X.T * W).dot(X)

В будущем мне пришлось бы запускать этот EM-GMM для данных еще большего размера (формы [2mil, 2500] и даже [2mil, 10k])
Что я могу сделать, чтобы сделать эти вычисления более эффективными в использовании памяти?

РЕДАКТИРОВАТЬ:

Я заметил, что рабочая инициализация использует pickle, поэтому X-матрицапревратился в ndarray, и работники не делятся им (что означает, что X-матрица дублируется для всех работников и заполняет мою оперативную память)

У меня есть представление о том, как ее решить, и я обновлю ее, если она исправлена.
Но если у кого-то есть хорошая идея, как с этим бороться, я буду благодарен.

Answer 1

Оказалось, что было 2 несвязанных проблемы, которые вызывали чрезмерное использование ОЗУ.

Во-первых, объект memmap полностью считывался с диска при обработке для многопроцессорных рабочих.
Это дублирование данных выделяло 6,7 ГБ ОЗУ для каждого рабочего.
Для решения этой проблемы яЯ создал общий RawArray и загрузил в него данные, и на каждом работнике я использовал np.frombuffer.

Во-вторых, и X.dot(Q), и (X.T * W) привели к беспорядочному выделению еще одной X-образной матрицы, то есть еще 6,7 ГБ ОЗУ
Я создал вариант ответа из этой темы: https://stackoverflow.com/a/21096605/5572523
Поскольку моя матрица невысока, я нарезал строки:

def _block_slices(dim_size, block_size):
    count = 0
    while True:
        yield slice(count, count + block_size, 1)
        count += block_size
        if count >= dim_size:
            raise StopIteration

И теперь я могу выполнять итерации по пакетам данных (также добавлено немного дополнительного ускорения при работе с weight = 0)

Я установил max_elements = 2 ** 27, потому что я использую float64, поэтому в результате получается матрица размером 1 ГБ (если я не ошибаюсь).

Итак, (X.T * W).dot(X) превратился в:

def weighted_outer_prod(X, W):
    n, d = X.shape

    max_rows = max(1, int(max_elements / d))
    sigma = np.zeros([d, d])
    for mm in _block_slices(n, max_rows):
        sigma += batch_weighted_outer_prod(X[mm, :], W[mm])
    return sigma

def batch_weighted_outer_prod(batch, W):
    nz = W > 0
    buff = np.empty(batch[nz].shape)
    np.multiply(batch[nz], W[nz, np.newaxis], out=buff)
    sigma = buff.T.dot(batch[nz])
    del(buff)
    return sigma

И np.sum(X.dot(Q) * X, axis=1) превратился в: (не обращая внимания на название функции)

def calc_l_k(X, Q):
    max_rows = max(1, int(max_elements / d))
    l_k = np.empty(n)
    for mm in _block_slices(n, max_rows):
        l_k[mm] = batch_l_k(X[mm, :], Q)

    return l_k


def batch_l_k(batch, Q):
    buff = np.empty(batch.shape)
    np.dot(batch, Q, out=buff)
    np.multiply(buff, batch, out=buff)
    l_k = -0.5 * np.sum(buff, axis=1)
    del(buff)
    return l_k

Теперь он работает с X формы [1 мил, 900], и я надеюсь, что это будетпо-прежнему работать с более высокими размерами.