大内存映射数组的高效点积

35

我正在处理一些比较大且密集的numpy浮点数组,这些数组目前驻留在PyTables的CArray中。我需要能够使用这些数组执行高效的点积计算,例如C = A.dot(B),其中A是一个巨大的(大约1E4 x 3E5 float32)内存映射数组,BC是较小的numpy数组,驻留在核心内存中。

目前我所做的是使用np.memmap将数据复制到内存映射的numpy数组中,然后直接在内存映射的数组上调用np.dot。虽然这种方法有效,但我怀疑标准的np.dot(或它调用的底层BLAS函数)在计算结果时所需的I/O操作数量可能不是很高效。

我在这篇综述文章中发现了一个有趣的例子。这是一种使用3个嵌套循环计算的朴素点积:

def naive_dot(A, B, C):
    for ii in xrange(n):
        for jj in xrange(n):
            C[ii,jj] = 0
            for kk in xrange(n):
                C[ii,jj] += A[ii,kk]*B[kk,jj]
    return C

需要 O(n^3) 的I/O操作来计算。

然而,通过以合适大小的块处理数组:

def block_dot(A, B, C, M):
    b = sqrt(M / 3)
    for ii in xrange(0, n, b):
        for jj in xrange(0, n, b):
            C[ii:ii+b,jj:jj+b] = 0
            for kk in xrange(0, n, b):
                C[ii:ii+b,jj:jj+b] += naive_dot(A[ii:ii+b,kk:kk+b], 
                                                B[kk:kk+b,jj:jj+b],
                                                C[ii:ii+b,jj:jj+b])
    return C

其中M是可容纳的最大元素数,I/O操作的数量减少到O(n^3 / sqrt(M))

np.dot和/或np.memmap有多聪明?调用np.dot会执行I/O有效的分块点积运算吗?np.memmap是否进行任何精细的缓存以改善此类操作的效率?

如果没有,是否存在某个现有库函数可以执行I/O有效的点积运算,还是我应该尝试自己实现它?

更新

我已经对一个手动实现的np.dot进行了一些基准测试,该实现对输入数组的块进行操作,并将其显式读入核心内存。这些数据至少部分回答了我的原始问题,因此我将其发布为答案。


1
SWAG:你有没有在Google Code和Cheese Factory上查看numexpr - Mark Mikofski
1
@MarkMikofski 谢谢,但那不是我正在寻找的东西 - 首先,因为我想对整个矩阵进行快速线性代数运算而不是逐元素操作,其次,因为在这种情况下我主要受到I/O限制而不是CPU限制。 - ali_m
4
当我说我是"I/O-bound"时,我的意思是主要拖慢我速度的因素是需要从硬盘读取数据到系统内存。如果限制因素是首先从硬盘读取数据,那么能够并行处理事物实际上并不能加速处理速度。 - ali_m
1
@J.F.Sebastian 我正在尝试实现用于近似大矩阵SVD的算法(http://arxiv.org/pdf/1007.5510.pdf)。我认为没有不使用矩阵乘法的方法来做到这一点。 - ali_m
1
@usethedeathstar 1) 我还没有尝试过np.einsum,因为我想不出它比np.dot更快的特定原因。对于计算两个在核心内存中的数组的点积,np.dot将比等效的np.einsum调用更快,因为它可以使用更高度优化的BLAS函数。在我的情况下,几乎没有区别,因为我受到I/O限制。2)不,正如我在描述中所说的,它们是密集的矩阵。 - ali_m
显示剩余8条评论
3个回答

26
我实现了一个函数,用于将`np.dot`应用于从内存映射数组中明确读入核心内存的块:
import numpy as np

def _block_slices(dim_size, block_size):
    """Generator that yields slice objects for indexing into 
    sequential blocks of an array along a particular axis
    """
    count = 0
    while True:
        yield slice(count, count + block_size, 1)
        count += block_size
        if count > dim_size:
            raise StopIteration

def blockwise_dot(A, B, max_elements=int(2**27), out=None):
    """
    Computes the dot product of two matrices in a block-wise fashion. 
    Only blocks of `A` with a maximum size of `max_elements` will be 
    processed simultaneously.
    """

    m,  n = A.shape
    n1, o = B.shape

    if n1 != n:
        raise ValueError('matrices are not aligned')

    if A.flags.f_contiguous:
        # prioritize processing as many columns of A as possible
        max_cols = max(1, max_elements / m)
        max_rows =  max_elements / max_cols

    else:
        # prioritize processing as many rows of A as possible
        max_rows = max(1, max_elements / n)
        max_cols =  max_elements / max_rows

    if out is None:
        out = np.empty((m, o), dtype=np.result_type(A, B))
    elif out.shape != (m, o):
        raise ValueError('output array has incorrect dimensions')

    for mm in _block_slices(m, max_rows):
        out[mm, :] = 0
        for nn in _block_slices(n, max_cols):
            A_block = A[mm, nn].copy()  # copy to force a read
            out[mm, :] += np.dot(A_block, B[nn, :])
            del A_block

    return out

我随后进行了一些基准测试,将我的blockwise_dot函数与直接应用于内存映射数组的普通np.dot函数进行比较(请参见以下基准测试脚本)。 我使用的是针对OpenBLAS v0.2.9.rc1(从源代码编译)的numpy 1.9.0.dev-205598b。该机器是运行Ubuntu 13.10的四核笔记本电脑,具有8GB RAM和SSD,并且已禁用交换文件。

结果

正如@Bi Rico所预测的那样,计算点积所需的时间与A的维数美妙地成为O(n)。在缓存的A块上操作比仅调用整个内存映射数组上的普通np.dot函数会大大提高性能:

enter image description here

它对正在处理的块的大小非常不敏感-在以1GB、2GB或4GB块处理数组所需的时间之间几乎没有什么区别。 我得出结论,无论np.memmap数组本地实现了什么缓存,它似乎都非常次优于计算点积。

进一步的问题

手动实现此缓存策略仍然有点麻烦,因为我的代码可能必须在具有不同物理内存量和潜在不同操作系统的机器上运行。出于这个原因,我仍然对是否有方法控制内存映射数组的缓存行为以提高np.dot的性能感兴趣。
当我运行基准测试时,我注意到了一些奇怪的内存处理行为-当我在整个A上调用np.dot时,我从未看到我的Python进程的常驻集大小超过约3.8GB,即使我有大约7.5GB的可用RAM。这导致我怀疑对于一个np.memmap数组,可能会对其允许占用的物理内存量施加一些限制-我之前认为它将使用操作系统允许它获取的任何RAM。在我的情况下,能够增加此限制可能非常有益。
是否有人对np.memmap数组的缓存行为有更深入的了解,可以帮助解释这一点?

基准测试脚本

def generate_random_mmarray(shape, fp, max_elements):
    A = np.memmap(fp, dtype=np.float32, mode='w+', shape=shape)
    max_rows = max(1, max_elements / shape[1])
    max_cols =  max_elements / max_rows
    for rr in _block_slices(shape[0], max_rows):
        for cc in _block_slices(shape[1], max_cols):
            A[rr, cc] = np.random.randn(*A[rr, cc].shape)
    return A

def run_bench(n_gigabytes=np.array([16]), max_block_gigabytes=6, reps=3,
              fpath='temp_array'):
    """
    time C = A * B, where A is a big (n, n) memory-mapped array, and B and C are
    (n, o) arrays resident in core memory
    """

    standard_times = []
    blockwise_times = []
    differences = []
    nbytes = n_gigabytes * 2 ** 30
    o = 64

    # float32 elements
    max_elements = int((max_block_gigabytes * 2 ** 30) / 4)

    for nb in nbytes:

        # float32 elements
        n = int(np.sqrt(nb / 4))

        with open(fpath, 'w+') as f:
            A = generate_random_mmarray((n, n), f, (max_elements / 2))
            B = np.random.randn(n, o).astype(np.float32)

            print "\n" + "-"*60
            print "A: %s\t(%i bytes)" %(A.shape, A.nbytes)
            print "B: %s\t\t(%i bytes)" %(B.shape, B.nbytes)

            best = np.inf
            for _ in xrange(reps):
                tic = time.time()
                res1 = np.dot(A, B)
                t = time.time() - tic
                best = min(best, t)
            print "Normal dot:\t%imin %.2fsec" %divmod(best, 60)
            standard_times.append(best)

            best = np.inf
            for _ in xrange(reps):
                tic = time.time()
                res2 = blockwise_dot(A, B, max_elements=max_elements)
                t = time.time() - tic
                best = min(best, t)
            print "Block-wise dot:\t%imin %.2fsec" %divmod(best, 60)
            blockwise_times.append(best)

            diff = np.linalg.norm(res1 - res2)
            print "L2 norm of difference:\t%g" %diff
            differences.append(diff)

        del A, B
        del res1, res2
        os.remove(fpath)

    return (np.array(standard_times), np.array(blockwise_times), 
            np.array(differences))

if __name__ == '__main__':
    n = np.logspace(2,5,4,base=2)
    standard_times, blockwise_times, differences = run_bench(
                                                    n_gigabytes=n,
                                                    max_block_gigabytes=4)

    np.savez('bench_results', standard_times=standard_times, 
             blockwise_times=blockwise_times, differences=differences)

请提交您的系统参数,以及确认您使用的Python、NumPy等软件包是否均为x64版本? - mrgloom
@mrgloom 一切都是x64。所有其他相关参数都在我的答案中描述。 - ali_m
你可能需要使用Strassen算法来减少切片的点积。但这会占用更多的内存。(我猜是指较小的切片) - Mehdi
@Mehdi 很高兴知道这对未来有用。不幸的是,目前我最受限制的是内存消耗。降低数值稳定性也可能是一个问题。我猜实际点积方面,要在速度上击败优化的BLAS函数会非常困难。 - ali_m

6

我认为numpy没有为memmap数组优化点积,如果你看一下矩阵乘法的代码,我在这里找到了它,你会发现函数MatrixProduct2(目前实现)按照c内存顺序计算结果矩阵的值:

op = PyArray_DATA(ret); os = PyArray_DESCR(ret)->elsize;
axis = PyArray_NDIM(ap1)-1;
it1 = (PyArrayIterObject *)
    PyArray_IterAllButAxis((PyObject *)ap1, &axis);
it2 = (PyArrayIterObject *)
    PyArray_IterAllButAxis((PyObject *)ap2, &matchDim);
NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ap2));
while (it1->index < it1->size) {
    while (it2->index < it2->size) {
        dot(it1->dataptr, is1, it2->dataptr, is2, op, l, ret);
        op += os;
        PyArray_ITER_NEXT(it2);
    }
    PyArray_ITER_NEXT(it1);
    PyArray_ITER_RESET(it2);
}

在上面的代码中,op是返回矩阵,dot是1d点积函数,it1it2是输入矩阵的迭代器。

话虽如此,看起来你的代码已经在做正确的事情。在这种情况下,最优性能实际上要比O(n^3/sprt(M))好得多,你可以将IO限制为只从磁盘读取A的每个项目一次,或者O(n)。 Memmap数组自然必须在幕后进行一些缓存操作,内层循环对it2进行操作,因此如果A按C顺序排列并且memmap高速缓存足够大,则您的代码可能已经在工作。您可以通过执行以下操作显式地强制缓存A的行:

def my_dot(A, B, C):

    for ii in xrange(n):
        A_ii = np.array(A[ii, :])
        C[ii, :] = A_ii.dot(B)

    return C

这让人感到放心 - 我想知道其他线性代数操作在何种程度上会与memmapped数组的缓存良好配合。您是否碰巧知道是否可以控制缓存大小?我从未找到过一份好的资源来解释memmap如何控制缓存和内存使用。 - ali_m
5
请注意,PyArray_MatrixProduct2仅在以下情况下由np.dot使用:不能调用BLAS(例如非BLAS兼容的内存顺序,非浮点数据类型,未安装BLAS库)。请参见此处 - pv.
基于它使用了我的4个核心,np.dot 在将一个memmapped的float32数组与一个非memmapped的float32数组相乘时似乎确实调用了BLAS,因此 PyArray_MatrixProduct2 可能不会被调用。 - ali_m

5
我建议您使用PyTables而不是numpy.memmap。此外,阅读他们关于压缩的演示,听起来很奇怪,但似乎序列"compress->transfer->uncompress"比仅传输未压缩的文件更快
同时,请使用带有MKL的np.dot。我不知道numexpr(pytables也似乎有类似的东西)如何用于矩阵乘法,但例如计算欧几里得范数,它是最快的方法(与numpy相比)。
请尝试对此示例代码进行基准测试:
import numpy as np
import tables
import time
n_row=1000
n_col=1000
n_batch=100
def test_hdf5_disk():
    rows = n_row
    cols = n_col
    batches = n_batch
    #settings for all hdf5 files
    atom = tables.Float32Atom()
    filters = tables.Filters(complevel=9, complib='blosc') # tune parameters
    Nchunk = 4*1024  # ?
    chunkshape = (Nchunk, Nchunk)
    chunk_multiple = 1
    block_size = chunk_multiple * Nchunk

    fileName_A = 'carray_A.h5'
    shape_A = (n_row*n_batch, n_col)  # predefined size
    h5f_A = tables.open_file(fileName_A, 'w')
    A = h5f_A.create_carray(h5f_A.root, 'CArray', atom, shape_A, chunkshape=chunkshape, filters=filters)
    for i in range(batches):
        data = np.random.rand(n_row, n_col)
        A[i*n_row:(i+1)*n_row]= data[:]
    rows = n_col
    cols = n_row
    batches = n_batch
    fileName_B = 'carray_B.h5'
    shape_B = (rows, cols*batches)  # predefined size
    h5f_B = tables.open_file(fileName_B, 'w')
    B = h5f_B.create_carray(h5f_B.root, 'CArray', atom, shape_B, chunkshape=chunkshape, filters=filters)
    sz= rows/batches
    for i in range(batches):
        data = np.random.rand(sz, cols*batches)
        B[i*sz:(i+1)*sz]= data[:]
    fileName_C = 'CArray_C.h5'
    shape = (A.shape[0], B.shape[1])
    h5f_C = tables.open_file(fileName_C, 'w')
    C = h5f_C.create_carray(h5f_C.root, 'CArray', atom, shape, chunkshape=chunkshape, filters=filters)
    sz= block_size
    t0= time.time()
    for i in range(0, A.shape[0], sz):
        for j in range(0, B.shape[1], sz):
            for k in range(0, A.shape[1], sz):
                C[i:i+sz,j:j+sz] += np.dot(A[i:i+sz,k:k+sz],B[k:k+sz,j:j+sz])
    print (time.time()-t0)
    h5f_A.close()
    h5f_B.close()
    h5f_C.close()

我不知道如何调整块大小和压缩率以适应当前计算机,因此我认为性能可能取决于参数。

另外请注意,示例代码中的所有矩阵都存储在磁盘上,如果其中一些矩阵将存储在RAM中,我认为速度会更快。

顺便说一下,我正在使用x32机器,使用numpy.memmap时矩阵大小存在一些限制(我不确定,但似乎视图大小只能达到~2Gb),而PyTables没有限制。


在PyTables数组上操作有一定的吸引力,部分原因是数据已经存储在PyTables数组中。然而,它们比numpy数组更棘手。我还必须对A的转置执行点积,由于它们缺少一个转置方法,这使得我的索引变得更加笨拙。最大的问题可能是选择适当的块形状,因为我还必须对A的单行/列执行操作,以及最好在方块上执行的点积。 - ali_m
PyTables数组是否比memmap数组更快,这完全取决于我的真实数据有多可压缩,因此我可以节省多少I/O带宽。不幸的是,我没有一个真实的数据集在我的本地机器上进行测试(正如我所说,它们相当大...),但我可以告诉你,使用我一直在使用的高斯合成数据时,使用PyTables CArrays而不是memmaps没有性能优势。这一点并不令人惊讶,因为随机数据从定义上来说是不可压缩的。我会在有机会时用真实数据进行一些基准测试。 - ali_m
以下是一些建议:https://groups.google.com/forum/#!topic/pytables-users/1jJhriRLDS4 - mrgloom
在那个帖子中,我认为Anthony Scopatz假设您的输入数组足够小,可以保存在内存中。当然,调用整个数组上的np.dot会更快,但我显然无法这样做。 - ali_m

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