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

28

做类似以下的事情

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

使用多个核心,运行效果良好。

然而,a 中的元素是 64 位浮点数(或在 32 位平台上为 32 位浮点数?),我想要乘以 8 位整数数组。尝试以下代码:

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

这导致点积在我的电脑上不能使用多个内核,因此运行速度慢了约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)

计算是多线程的。现在,对于float32,blas.sgemm与np.dot完全具有相同的时间运行,但对于非floats,它将所有内容转换为float32并输出浮点数,而这是np.dot所不做的。(此外,b现在以F_CONTIGUOUS顺序排列,这是一个较小的问题)。
因此,如果我想进行整数矩阵乘法,我必须执行以下操作之一:
1.使用NumPy的极慢的np.dot,并且很高兴保留8位整数。 2.使用SciPy的sgemm并使用4倍内存。 3.使用Numpy的np.float16,并仅使用2倍内存,但要注意np.dot在float16数组上比在float32数组上慢得多,特别是int8。 4.查找优化的库以进行多线程整数矩阵乘法(实际上,Mathematica可以做到这一点,但我更喜欢Python解决方案),最好支持1位数组,尽管8位数组也可以...(我实际上旨在对Z / 2Z上的矩阵进行乘法,我知道我可以用Sage做到这一点,它非常像Python,但是,还有严格的Python吗?)
我能按照选项4吗?是否存在这样的库?
免责声明:我实际上正在运行NumPy + MKL,但我已经尝试过vanilly NumPy上的类似测试,并获得了类似的结果。

1
关于您的选项n°4,也许您可以看看PyCuda或者Theano?它们允许在GPU上进行大型操作(与numpy接口方便),性能相当不错。 - mgc
1
作为第四个选项的可能答案,https://bitbucket.org/malb/m4ri 看起来很有趣。 "M4RI 是一个用于在F2上进行稠密矩阵快速算术的库"。我想这已经是Sage正在使用的,但我看不到任何理由为什么你不能直接从Python中使用它,只需适当的Cython包装器。(实际上,您可能已经能够在Sage源代码中找到这样的包装器。) - Mark Dickinson
1
还没有人提到 numpy.einsum,但这可能是一个不错的选择。 - user2379410
1
请注意,如果您想避免整数溢出,则需要将结果转换为大一些的类型。 如果每个元素只是0或1,则需要一个能够保存至少“n”值的整数格式以保证不会溢出。 对于您的示例,其中“n = 10000”,(u)int16应该足够了。 您的实际矩阵很稀疏吗? 如果是的话,最好使用“scipy.sparse.csr_matrix”。 - ali_m
1
你能提供一些关于你试图解决的整体问题的更多背景吗?将大整数矩阵相乘在某种程度上是一件不寻常的事情。了解这些矩阵的属性会特别有用。这些值是否总是0或1?如果它们可以更大,那么你可能最终会受到使用uint64表示的最大整数的限制。这些矩阵是如何生成的?它们是否具有任何特殊结构(例如对称性、块、带等)? - ali_m
显示剩余9条评论
2个回答

8
请注意,虽然这个答案已经有点老了,但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))
    

通过这个实现,我获得了大约4倍的加速,这是我的机器上物理核心数。

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

它运行了!这是O(n**3)矩阵乘积,确切地做了n**2个点积,对吗? - étale-cohomology
它将矩阵乘积分成多个较小的矩阵乘积。在极端情况下,这可以是向量点积。 - MB-F
当类型为float时,pardot比np.dot慢:同时运行4个作业、同时运行8个作业、pardot:0.13秒、np.dot:0.07秒 - kory
当数据集的大小增加到10倍时,情况变得更糟:
pardot: 耗时1212.89秒
np.dot: 耗时73.11秒
- kory
@kory 这是可以预料的。请使用 np.dot 进行浮点数乘法。 - MB-F
@kazemakase 是的,找到两种情况的解决方案会很好,我只是指出如果有人也用这个来处理浮点数值时会有多糟糕。谢谢! - kory

2
"为什么执行浮点数矩阵乘法比整数矩阵乘法快?"解释了为什么整数运算很慢:首先,CPU有高吞吐量的浮点数管道。其次,BLAS没有整数类型。

解决方法:将矩阵转换为float32值可以获得大幅提速。在2015年的MacBook Pro上,速度提升了90倍!(使用float64只有一半的效果。)

"
import numpy as np
import time

def timeit(callable):
    start = time.time()
    callable()
    end = time.time()
    return end - start

a = np.random.random_integers(0, 9, size=(1000, 1000)).astype(np.int8)

timeit(lambda: a.dot(a))  # ≈0.9 sec
timeit(lambda: a.astype(np.float32).dot(a.astype(np.float32)).astype(np.int8) )  # ≈0.01 sec

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接