双重 Python 循环的 NumPy 向量化

Question

双重 Python 循环的 NumPy 向量化

4

V是一个(n,p)的numpy数组，通常维度为n~10，p~20000。

我现在拥有的代码看起来像这样：

A = np.zeros(p)
for i in xrange(n):
    for j in xrange(i+1):
        A += F[i,j] * V[i,:] * V[j,:]

我该如何改写这段代码以避免使用双重Python循环？

- user1984528

1

F的形状是什么？它是(n, n)还是(n, n, p)？ - wflynny

你考虑过用C语言（因为它们很简单）编写这些for循环，并使用scipy.weave.blitz吗？ - usethedeathstar

如果一些Python代码“足够好”，我宁愿暂时避免使用Cython或Weave.blitz。 - user1984528

3个回答

7

这个问题的难点在于你只想对 j <= i 的元素求和。如果不是这样，你可以按照以下方法操作：

M = (V.reshape(n, 1, p) * V.reshape(1, n, p)) * F.reshape(n, n, 1)
A = M.sum(0).sum(0)

如果F是对称的（如果F[i,j] == F[j,i]），那么您可以如下利用M的对称性：

D = M[range(n), range(n)].sum(0)
A = (M.sum(0).sum(0) - D) / 2.0 + D

话虽如此，这并不是一个很好的向量化候选项，因为你有 n << p，所以你的 for 循环对该计算的速度影响不大。

编辑：像下面的 Bill 所说，你可以先确保你不想使用的 F 元素设置为零，然后 M.sum(0).sum(0) 的结果就是你想要的。

- Isaac

3

F[np.triu_indices(n, 1)] = 0的作用是什么？这样，F的上半部分（偏移了1）为零，不会对总和产生贡献。 - wflynny

此外，a.sum(0).sum(0) 可以写成 np.sum(a, (0,1)) 或者 a.sum((0,1))，虽然速度上并没有提升，但我认为更易于阅读。 - askewchan

就像我说的，这段代码不是向量化的好选择。 - Isaac

1

这句话可以翻译为：“表达式可以写成”。同时需要保留HTML标签，不做解释。

$A_p = \sum_{ij}F_{ij}V_{ip}V_{jp}$

因此，您可以使用np.newaxis构造来进行如下求和：

na = np.newaxis
X = (np.tri(n)*F)[:,:,na]*V[:,na,:]*V[na,:,:]
X.sum(axis=1).sum(axis=0)

这里构建了一个3D数组X [i，j，p]，然后对前两个轴求和，得到一个1D数组A [p]。此外，根据问题的限制，还将F乘以三角矩阵来限制求和。

- flonk

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Jaime · Accepted Answer

虽然Isaac的答案看起来很有前途，因为它消除了这两个嵌套的for循环，但你需要创建一个中间数组M，其大小是原始数组V的n倍。Python for循环不便宜，但内存访问也不是免费的：

n = 10
p = 20000
V = np.random.rand(n, p)
F = np.random.rand(n, n)

def op_code(V, F):
    n, p = V.shape
    A = np.zeros(p)
    for i in xrange(n):
        for j in xrange(i+1):
            A += F[i,j] * V[i,:] * V[j,:]
    return A

def isaac_code(V, F):
    n, p = V.shape
    F = F.copy()
    F[np.triu_indices(n, 1)] = 0
    M = (V.reshape(n, 1, p) * V.reshape(1, n, p)) * F.reshape(n, n, 1)
    return M.sum((0, 1))

如果你现在对这两个进行测试骑行：

In [20]: np.allclose(isaac_code(V, F), op_code(V, F))
Out[20]: True

In [21]: %timeit op_code(V, F)
100 loops, best of 3: 3.18 ms per loop

In [22]: %timeit isaac_code(V, F)
10 loops, best of 3: 24.3 ms per loop

因此，删除for循环会导致8倍减速。这不是一件好事...此时，您甚至可能要考虑是否需要进一步优化约需3毫秒才能评估的函数。如果需要，可以通过使用np.einsum实现微小的改进：

def einsum_code(V, F):
    n, p = V.shape
    F = F.copy()
    F[np.triu_indices(n, 1)] = 0
    return np.einsum('ij,ik,jk->k', F, V, V)

现在：

In [23]: np.allclose(einsum_code(V, F), op_code(V, F))
Out[23]: True

In [24]: %timeit einsum_code(V, F)
100 loops, best of 3: 2.53 ms per loop

所以大致上可以提高20%的速度，但是引入了可能不如for循环容易读懂的代码。我认为并不值得...