向量化这个 for 循环

Question

我有一个需要矢量化的 for 循环。下面的代码可以工作，但会花费很多时间（这是一个简化的示例，完整版本在 col_ids 中将有大约 1e6 行）。有人能告诉我如何矢量化这段代码以摆脱循环吗？如果重要，col_ids 是固定的（每次代码为运行时都相同），而 values 会改变。

values = np.array([1.5, 2, 2.3])
col_ids = np.array([[0,0,0,0], [0,0,0,1], [0,0,1,1]])
result = np.zeros((4,3))
for idx, col_idx in enumerate(col_ids):
    result[np.arange(4),col_idx] += values[idx]

结果：

[[5.8 0.  0. ]
 [5.8 0.  0. ]
 [3.5 2.3 0. ]
 [1.5 4.3 0. ]]

更新：我正在添加第二个示例，因为我的第一个示例的维度存在一些歧义。只有 values 和 col_ids 被更新，其他的都和第一个例子一样。（我保留第一个，因为答案中提到了这个）

values = np.array([1.5, 2, 5, 20, 50])
col_ids = np.array([[0,0,0,0], [0,0,0,1], [0,0,1,1], [0,0,1,2], [0,1,2,2]])

结果：

[[78.5  0.   0. ]
 [28.5 50.   0. ]
 [ 3.5 25.  50. ]
 [ 1.5  7.  70. ]]

所以 result 是 m x n，col_ids 是 k x m，值的长度是 k。 m和n都很小（m=4，n=3），k很大（完整例子大约1e6）

Answer 1

您可以使用 np.add.at 解决此问题。但是，据我所知，此函数不支持二维数组，因此您需要展平数组，计算一维展平索引，然后调用函数：

n, m = result.shape
result = np.zeros((4,3))
indices = np.tile(np.arange(0, n*m, m), col_ids.shape[0]) + col_ids.ravel()
np.add.at(result.ravel(), indices, np.repeat(values, n)) # In-place
print(result)

Answer 2

您可以对循环进行矢量化，但创建额外的中间数组对于较大的数据要慢得多（从 result 开始，形状为 (50,50)）

import numpy as np

values = np.array([1.5, 2, 2.3])
col_ids = np.array([[0,0,0,0], [0,0,0,1], [0,0,1,1]])

(np.equal.outer(col_ids, np.arange(len(values))) * values[:,None,None]).sum(0)

# for a fixed result shape (4,3)
# (np.equal.outer(col_ids, np.arange(3)) * values[:,None,None]).sum(0)

输出

array([[5.8, 0. , 0. ],
       [5.8, 0. , 0. ],
       [3.5, 2.3, 0. ],
       [1.5, 4.3, 0. ]])

我能找到的唯一可靠且更快的解决方案是 numba（使用 version 0.55.1）。我认为此实现会受益于并行执行，但我无法在 2 核 colab 实例上获得任何加速。

import numba as nb

@nb.njit(parallel=False) # Try parallel=True for multi-threaded execution, no speed up in my benchmarks 
def fill(val, ids):
    res = np.zeros(ids.shape[::-1])
    for i in nb.prange(len(res)):
        for j in range(res.shape[1]):
            res[i, ids[j,i]] += val[j]
    return res

fill(values, col_ids)

输出

array([[5.8, 0. , 0. ],
       [5.8, 0. , 0. ],
       [3.5, 2.3, 0. ],
       [1.5, 4.3, 0. ]])

对于固定的结果形状 (4,3) 和合适的输入。

@nb.njit(boundscheck=True) # ~1.25x slower, but much safer
def fill(val, ids):
    res = np.zeros((4,3))
    for i in nb.prange(ids.shape[0]):              
        for j in range(ids.shape[1]):   
            res[j, ids[i,j]] += val[i]
    return res

fill(values, col_ids)

更新示例数据的输出

array([[78.5,  0. ,  0. ],
       [28.5, 50. ,  0. ],
       [ 3.5, 25. , 50. ],
       [ 1.5,  7. , 70. ]])

向量化这个 for 循环

Vectorize this for loop

python

numpy

vectorization