python - NumPy/SciPy 中的多线程整数矩阵乘法

Question

做类似的事情

import numpy as np
a = np.random.rand(10**4, 10**4)
b = np.dot(a, a)

使用多核，运行良好。

a 中的元素是 64 位 float (或 32 位平台中的 32 位？)，我想乘以 8 位整数数组。不过，请尝试以下方法:

a = np.random.randint(2, size=(n, n)).astype(np.int8)

导致点积不使用多个内核，因此在我的 PC 上运行速度慢了约 1000 倍。

array: np.random.randint(2, size=shape).astype(dtype)

dtype    shape          %time (average)

float32 (2000, 2000)    62.5 ms
float32 (3000, 3000)    219 ms
float32 (4000, 4000)    328 ms
float32 (10000, 10000)  4.09 s

int8    (2000, 2000)    13 seconds
int8    (3000, 3000)    3min 26s
int8    (4000, 4000)    12min 20s
int8    (10000, 10000)  It didn't finish in 6 hours

float16 (2000, 2000)    2min 25s
float16 (3000, 3000)    Not tested
float16 (4000, 4000)    Not tested
float16 (10000, 10000)  Not tested

我知道 NumPy 使用 BLAS，它不支持整数，但如果我使用 SciPy BLAS 包装器，即。

import scipy.linalg.blas as blas
a = np.random.randint(2, size=(n, n)).astype(np.int8)
b = blas.sgemm(alpha=1.0, a=a, b=a)

计算是多线程的。现在，blas.sgemm 的运行时间与 float32 的 np.dot 完全相同，但对于非 float ，它将所有内容都转换为 float32 和输出 float ，这是 np.dot 不做的。 (此外，b 现在是 F_CONTIGUOUS 顺序，这是一个较小的问题)。

所以，如果我想进行整数矩阵乘法，我必须执行以下操作之一:

1. 使用 NumPy 非常缓慢的 np.dot，很高兴我能保留 8 位整数。
2. 使用 SciPy 的 sgemm 并使用 4x 内存。
3. 使用 Numpy 的 np.float16 并且只使用 2x 内存，但需要注意的是 np.dot 在 float16 数组上比在 float32 数组上慢得多，比int8.
4. 为多线程整数矩阵乘法找到一个优化的库(实际上，Mathematica 可以做到这一点，但我更喜欢 Python 解决方案)，理想情况下支持 1 位数组，尽管 8 位数组也很好......(我实际上的目标是在有限域 Z/2Z 上进行矩阵乘法，并且我知道我可以使用 Sage 来做到这一点，这很 Pythonic，但是，再次，有什么严格意义上的 Python 吗？)

我可以遵循选项 4 吗？有这样的图书馆吗？

免责声明:我实际上是在运行 NumPy + MKL，但我在 vanilly NumPy 上尝试了类似的测试，结果类似。

Answer 1

请注意，虽然这个答案变得陈旧，但 numpy 可能会获得优化的整数支持。请验证此答案在您的设置中是否仍然可以更快地工作。

• 选项 5 - 推出自定义解决方案: 将矩阵产品划分为几个子产品并并行执行。使用标准 Python 模块可以相对容易地实现这一点。子产品是用 numpy.dot 计算的，它会释放全局解释器锁。因此，可以使用 threads它们相对轻量级，可以从主线程访问数组以提高内存效率。

实现:

import numpy as np
from numpy.testing import assert_array_equal
import threading
from time import time


def blockshaped(arr, nrows, ncols):
    """
    Return an array of shape (nrows, ncols, n, m) where
    n * nrows, m * ncols = arr.shape.
    This should be a view of the original array.
    """
    h, w = arr.shape
    n, m = h // nrows, w // ncols
    return arr.reshape(nrows, n, ncols, m).swapaxes(1, 2)


def do_dot(a, b, out):
    #np.dot(a, b, out)  # does not work. maybe because out is not C-contiguous?
    out[:] = np.dot(a, b)  # less efficient because the output is stored in a temporary array?


def pardot(a, b, nblocks, mblocks, dot_func=do_dot):
    """
    Return the matrix product a * b.
    The product is split into nblocks * mblocks partitions that are performed
    in parallel threads.
    """
    n_jobs = nblocks * mblocks
    print('running {} jobs in parallel'.format(n_jobs))

    out = np.empty((a.shape[0], b.shape[1]), dtype=a.dtype)

    out_blocks = blockshaped(out, nblocks, mblocks)
    a_blocks = blockshaped(a, nblocks, 1)
    b_blocks = blockshaped(b, 1, mblocks)

    threads = []
    for i in range(nblocks):
        for j in range(mblocks):
            th = threading.Thread(target=dot_func, 
                                  args=(a_blocks[i, 0, :, :], 
                                        b_blocks[0, j, :, :], 
                                        out_blocks[i, j, :, :]))
            th.start()
            threads.append(th)

    for th in threads:
        th.join()

    return out


if __name__ == '__main__':
    a = np.ones((4, 3), dtype=int)
    b = np.arange(18, dtype=int).reshape(3, 6)
    assert_array_equal(pardot(a, b, 2, 2), np.dot(a, b))

    a = np.random.randn(1500, 1500).astype(int)

    start = time()
    pardot(a, a, 2, 4)
    time_par = time() - start
    print('pardot: {:.2f} seconds taken'.format(time_par))

    start = time()
    np.dot(a, a)
    time_dot = time() - start
    print('np.dot: {:.2f} seconds taken'.format(time_dot))
    

通过这个实现，我得到了大约 x4 的加速，这是我机器中的物理内核数:

running 8 jobs in parallel
pardot: 5.45 seconds taken
np.dot: 22.30 seconds taken