Как повысить производительность функций NumPy, связанных с BLAS?

Question

У меня есть функция, которая возвращает квадрат нормы невязки большой системы линейных уравнений.

In [1]: import numpy as np                                                      

In [2]: A = np.random.rand(3600000, 200)                                        

In [3]: b = np.random.rand(3600000)                                             

In [4]: def f(x): 
   ...:     global A
   ...:     global b
   ...:     return np.linalg.norm(A.dot(x) - b)**2                                       

Теперь у меня есть алгоритм, в котором функция должна вычисляться несколько раз. Однако из-за размера системы уравнений каждый вызов функции на определенном x требует много времени.

In [5]: import time                                                             

In [6]: def f(x): 
   ...:     global A 
   ...:     global b 
   ...:     start = time.time() 
   ...:     res = np.linalg.norm(A.dot(x) - b)**2 
   ...:     end = time.time() 
   ...:     return res, end - start 

In [7]: test = np.random.rand(200)                                             

In [8]: f(test)                                                                
Out[8]: (8820030785.528395, 7.467242956161499)

Мой вопрос:

Есть ли есть ли возможности сократить время вызова такой функции?

Я думал о замене np.linalg.norm(A.dot(x) - b)**2 более эффективным выражением, но не знаю, как это могло бы выглядеть.

Техническая информация. Приведенный выше код был написан на компьютере с

• macOS Catalina версии 10.15.5
• 2,3 ГГц Dual-Core Intel Core i5 (ускорение Turbo Boost до 3,6 ГГц) с 64 МБ eDRAM
• 8 ГБ 2133 МГц LPDDR3 RAM (встроенная)

•   Memory:
  
    Memory Slots:
  
     ECC: Disabled
     Upgradeable Memory: No
  
       BANK 0/DIMM0:
  
         Size: 4 GB
         Type: LPDDR3
         Speed: 2133 MHz
         Status: OK (...)
  
       BANK 1/DIMM0:
  
         Size: 4 GB
         Type: LPDDR3
         Speed: 2133 MHz
         Status: OK (...)
  

Результат np.show_config():

blas_mkl_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack']
    library_dirs = ['/Users/me/miniconda3/envs/magpy/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/Users/me/miniconda3/envs/magpy/include']
blas_opt_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack', 'blas', 'cblas', 'lapack']
    library_dirs = ['/Users/me/miniconda3/envs/magpy/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/Users/me/miniconda3/envs/magpy/include']
lapack_mkl_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack']
    library_dirs = ['/Users/me/miniconda3/envs/magpy/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/Users/me/miniconda3/envs/magpy/include']
lapack_opt_info:
    libraries = ['blas', 'cblas', 'lapack', 'pthread', 'blas', 'cblas', 'lapack', 'blas', 'cblas', 'lapack']
    library_dirs = ['/Users/me/miniconda3/envs/magpy/lib']
    define_macros = [('SCIPY_MKL_H', None), ('HAVE_CBLAS', None)]
    include_dirs = ['/Users/me/miniconda3/envs/magpy/include']

Answer 1

Проблема с производительностью, по-видимому, связана с медленной реализацией BLAS по умолчанию .

Реализация BLAS по умолчанию, используемая на вашем компьютере, по-видимому, Intel MKL , которая является обычно довольно быстро, но неожиданно медленно на вашей машине. Действительно, исходя из предоставленной информации об оборудовании, время выполнения должно составлять около 170-200 мс, а не 7,5 секунд.

Вы можете попробовать переключиться на другую реализацию BLAS , такую ​​как OpenBLAS, Apple Accelerate или BLIS. Вы можете найти информацию о том, как это сделать здесь и здесь .

Если переход на другую реализацию BLAS не решает проблему, вы используете следующую резервная реализация Numba :

@njit(parallel=True)
def customMathOp(A, x, b):
    squareNorm = 0.0
    for i in prange(A.shape[0]):
        s = 0.0
        for j in range(A.shape[1]):
            s += A[i,j] * x[j]
        squareNorm += (s - b[i]) ** 2
    return squareNorm

def f(x):
    global A
    global b
    start = time.time()
    res = customMathOp(A, x, b)
    end = time.time()
    return res, end - start

Этот код не так хорош, как использование numpy функций, основанных на быстрой реализации BLAS, но он должен быть относительно быстрым (обратите внимание, что первый вызов f будет немного медленным, так как время компиляции включено).

Обратите внимание, что использование типа np.float32 для массивов может ускорить выполнение в 2 раза, хотя результат также должен быть менее точным .

Answer 2

В вашем случае np.linalg.norm это просто

np.sqrt(dot(x,x))

Так что вам может быть лучше сделать:

temp = np.dot(A,x) - b         # temp = A@x-b
return np.dot(temp, temp)      # return temp@temp

Пропуск ненужных sqrt / square. Но по сравнению с исходным A@x это может быть незначительной проблемой.

На довольно простом компьютере Linux 4 ГБ ваш тестовый пример дает мне (при создании A)

MemoryError: Unable to allocate 5.36 GiB for an array with shape (3600000, 200) and data type float64

Хотя у вас, по-видимому, достаточно памяти, вы можете выйти за эту границу. В других SO мы видели, что dot/@ с очень большими массивами замедляется из-за проблем с управлением памятью. Часто люди набирают скорость, выполняя некую обработку «фрагментов». Это легко, если вы выполняете matmul с 3D-пакетом. Менее очевидно это с вашим нормальным корпусом.

Урезание A размер на 10:

In [423]: A.shape                                                                                    
Out[423]: (360000, 200)
In [424]: temp = A@x-b; res = temp@temp                                                              
In [425]: res                                                                                        
Out[425]: 938613433.9717302
In [426]: np.linalg.norm(A.dot(x)-b)**2                                                              
Out[426]: 938613433.9717301

Не сильно отличается в разы:

In [428]: timeit temp = A@x-b; res = temp@temp                                                       
85 ms ± 529 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [429]: timeit np.linalg.norm(A.dot(x)-b)**2                                                       
86.1 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

На самом деле это A.dot(x), который доминирует во времени; остальное незначительно.

Удвоение размера A, примерно вдвое больше (диапазон 175-180).

Я не специалист по библиотекам, но я верю MKL - один из самых быстрых вариантов, которого у меня нет (но у вас, похоже, есть).