点乘不同 dtype 的大密集矩阵（float x boolean）

Question

我正在乘以 2 个矩阵，A.dot(B)，其中：

A = 1 x n 矩阵，dtype float

B = n x n 矩阵，dtype 布尔值

我正在为较大的 n 执行此计算，并且运行内存很快就会耗尽（大约 n=14000 失败）。 A和B是密集的。

看来是因为numpy在执行矩阵乘法之前将B转换为dtype float，因此产生了巨大的内存成本。事实上，%timeit 表明将 B 转换为浮点数比执行乘法花费的时间更多。

有办法解决这个问题吗？这里的重点是减少内存尖峰/浮点转换，同时仍然允许常见的矩阵功能（矩阵加法/乘法）。

以下是基准测试解决方案的可重现数据：

np.random.seed(999)
n = 30000
A = np.random.random(n)
B = np.where(np.random.random((n, n)) > 0.5, True, False)

Answer 1

您可以节省 space 和使用 np.packbits 将布尔数组压缩到位域的时间，然后在行上 np.bincount 以同时计算 8 个标量积的块。

import numpy as np

def setup_data(M, N):
    return {'B': np.random.randint(0, 2, (M, N), dtype=bool),
            'A': np.random.random((M,))}

def f_vecmat_mult(A, B, decode=np.array(np.unravel_index(np.arange(256), 8*(2,)))):
    M, N = B.shape
    out = [(decode * np.bincount(row, A, minlength=256)).sum(axis=1) for row in np.packbits(B, axis=1).T]
    if N & 7:
        out[-1] = out[-1][:N & 7]
    return np.concatenate(out)

def f_direct(A, B):
    return A @ B

import types
from timeit import timeit

for M, N in [(99, 80), (999, 777), (9999, 7777), (30000, 30000)]:
    data = setup_data(M, N)
    ref = f_vecmat_mult(**data)
    print(f'M, N = {M}, {N}')
    for name, func in list(globals().items()):
        if not name.startswith('f_') or not isinstance(func, types.FunctionType):
            continue
        try:
            assert np.allclose(ref, func(**data))
            print("{:16s}{:16.8f} ms".format(name[2:], timeit(
                'f(**data)', globals={'f':func, 'data':data}, number=100)*10))
        except:
            print("{:16s} apparently failed".format(name[2:]))

示例输出：

M, N = 99, 80
vecmat_mult           0.12248290 ms
direct                0.03647798 ms
M, N = 999, 777
vecmat_mult           1.67854790 ms
direct                5.68286091 ms
M, N = 9999, 7777
vecmat_mult          68.74523309 ms
direct              571.34140913 ms
M, N = 30000, 30000
vecmat_mult        1345.18991556 ms
direct           apparently failed

点乘不同 dtype 的大密集矩阵（float x boolean）

Dot-multiplying large dense matrices of different dtype (float x boolean)

python

numpy

matrix

type-conversion

sparse-matrix