Умножение больших матриц с DASK

Question

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

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) не очень (и поэтому нелегко поддерживать).

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

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

Любые советы о том, как обойти это, будут высоко оценены.

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

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

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

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

Редактировать2: Это очень упрощенный пример того, что я пытался сделать. Проблема заключается, я считаю, в рекурсиях, связанных с определением столбцов в 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 делает большую разницу в скорости и удобочитаемости.

Answer 1

Ваш выход очень большой.

>>> A.T.dot(A).shape 
(10000000, 10000000) 

Возможно, вы намеревались вычислить это с транспозицией в другом направлении?

>>> A.dot(A.T).shape 
(2000, 2000) 

Это все еще занимает некоторое время (это большое вычисление), но оно завершается.