python - 向量化这个 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

您可以对循环进行矢量化处理，但是对于较大的数据(从 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. ]])