Умножение больших матриц с помощью dask ⇐ Python

[Anonymous]

Я работаю над проектом, который по сути сводится к решению матричного уравнения

Код: Выделить всё

A.dot(x) = d

где A — матрица размером примерно 10 000 000 на 2000 (в конечном итоге я бы хотел увеличить ее в обоих направлениях).

Код: Выделить всё

A

явно не помещается в памяти, поэтому его необходимо распараллелить. Вместо этого я делаю это, решая A.T.dot(A).dot(x) = A.T.dot(d). A.T будет иметь размеры 2000 на 2000. Его можно вычислить, разделив A и d на фрагменты A_i и d_i, вычислив по строкам A_i.T.dot(A_i) и A_i.T.dot(d_i) и суммировав их. Идеально подходит для распараллеливания. Мне удалось реализовать это с помощью модуля многопроцессорной обработки, но 1) его трудно масштабировать дальше (увеличивая A в обоих измерениях) из-за использования памяти, и 2) это некрасиво (и, следовательно, его нелегко поддерживать).

Dask кажется очень многообещающей библиотекой для решения обеих этих проблем, и я предпринял несколько попыток. Мою матрицу A сложно вычислить: она основана примерно на 15 различных массивах (размер которых равен количеству строк в A), и некоторые из них используются в итеративном алгоритме для вычисления связанной функции Лежандра. Когда куски небольшие (10000 строк), построение графа задачи занимает очень много времени и требует много памяти (увеличение памяти совпадает с вызовом итерационного алгоритма). Когда фрагменты больше (50 000 строк), потребление памяти перед вычислениями намного меньше, но оно быстро исчерпывается при вычислении A.T.dot(A). Я пробовал использовать кэш.Chest, но он значительно замедляет вычисления.

График задачи должен быть очень большим и сложным — вызов A._visualize() приводит к сбою. С более простыми матрицами A это можно сделать напрямую (см. ответ @MRocklin). Есть ли способ упростить это?

Буду очень признателен за любые советы о том, как обойти эту проблему.

В качестве игрушечного примера я попробовал

Код: Выделить всё

A = da.ones((2e3, 1e7), chunks = (2e3, 1e3))
(A.T.dot(A)).compute()

Это также не удалось: была использована вся память при активном только одном ядре. При чанках = (2e3, 1e5) все ядра запускаются почти сразу, но MemoryError появляется в течение 1 секунды (у меня на нынешнем компьютере 15 ГБ). chunks = (2e3, 1e4) был более многообещающим, но в итоге он также потреблял всю память.

Редактировать:
Я зачеркнул тест с примером игрушки, потому что размеры были неправильными, и исправил размеры в остальных случаях. Как говорит @MRocklin, он работает с правильными размерами. Я добавил вопрос, который, как мне кажется, больше соответствует моей проблеме.

Edit2:
Это очень упрощенный пример того, что я пытался сделать. Я считаю, что проблема в рекурсиях, связанных с определением столбцов в A.

Код: Выделить всё

import dask.array as da

N = 1e6
M = 500

x = da.random.random((N, 1), chunks = 5*M)

# my actual
A_dict = {0:x}
for i in range(1, M):
A_dict[i] = 2*A_dict[i-1]
A = da.hstack(tuple(A_dict.values()))
A = A.rechunk((M*5, M))
ATA = A.T.dot(A)

Похоже, это приводит к очень сложному графу задач, который занимает много памяти еще до начала вычислений.

Теперь я решил эту проблему, поместив рекурсию в функцию с массивами numpy и более или менее A = x.map_blocks(...).

Во-вторых, как только у меня есть граф задач матрицы A, непосредственное вычисление A.T.dot(A) действительно приводит к некоторым проблемам с памятью (использование памяти не очень стабильно). Поэтому я явно вычисляю его по частям и суммирую результаты. Даже с этими обходными путями dask существенно повышает скорость и читабельность.

Подробнее здесь: https://stackoverflow.com/questions/353 ... -with-dask