python - 如何加速大量内积

Question

我想遍历由 -1 和 1 组成的所有 2^32 行，并对每一行执行内积运算。有没有办法加快下面的代码？

import itertools
import operator

n = 16

survivors = []
for row in itertools.product([-1,1], repeat = 2*n):
    if (sum(map(operator.mul, row[0:n], row[n:2*n])) == 0):
        survivors.append(row)
print len(survivors)

Answer 1

这有点棘手，但是由 -1 和 1 组成的向量的内积可以转换为 XOR 运算和非零计数由 0 和 1 组成的向量项。当然，0 和 1 的 32 项向量的最佳容器是 uint32。下面的代码与您建议的相同，但是以 2**m 元素的 block 矢量化运行它以提高速度

def bitwise_inner(n, m=12):
    bitmask = (1 << n) - 1
    m = min(m, 2*n)
    result = []
    chunk = 2**m
    for j in xrange(2**(2*n-m)):
        start = j * chunk
        items = np.arange(start, start+chunk, dtype=np.uint32)
        op = (items >> n) ^ (items & bitmask)
        # Keep only items with same number of 0s and 1s in the first n bits of op
        for k in xrange(n//2 - 1):
            # This removes a 1 from the binary representation of each number
            op &= op - 1
            mask = op != 0
            items = items[mask]
            op = op[mask]
        op &= op - 1
        result.append(items[op == 0])
    return sum([len(j) for j in result])

首先让我们检查您的代码的正确性:

def python_inner(n):
    survivors = []
    for row in itertools.product([-1,1], repeat = 2*n):
        if (sum(map(operator.mul, row[0:n], row[n:2*n])) == 0):
            survivors.append(row)
    return len(survivors)

In [2]: python_inner(8)
Out[2]: 17920

In [3]: bitwise_inner(8)
Out[3]: 17920

然后是一些时间:

In [4]: %timeit python_inner(8)
1 loops, best of 3: 1.08 s per loop

In [5]: %timeit bitwise_inner(8)
100 loops, best of 3: 3.35 ms per loop

In [6]: %timeit bitwise_inner(12)
1 loops, best of 3: 1.07 s per loop

计算 n = 16 的所有值仍将花费大量时间，但这仍然快了两个数量级。