python矢量化：如何提高4層循環的效率

Question

我正在嘗試使用 Gibbs 采樣實現 LDA，在更新每個主題比例的步驟中，我有一個 4 層循環，它運行非常慢，我不知道如何提高這段代碼的效率。 我現在擁有的代碼如下：

N_W 是單詞數，N_D 是文檔數，Z[i,j] 是主題分配（1 到 K 個可能的分配），X[i,j] 是第 j 個單詞的計數第 i 個文檔 Beta[k,:] 的維度為 [K, N_W]。

更新如下：

for k in range(K): # iteratively for each topic update
    n_k = np.zeros(N_W) # vocab size

    for w in range(N_W):
        for i in range(N_D):
            for j in range(N_W): 
                # counting number of times a word is assigned to a topic
                n_k[w] += (X[i,j] == w) and (Z[i,j] == k) 

    # update
    Beta[k,:] = np.random.dirichlet(gamma + n_k)

Answer 1

您可以使用邏輯函數擺脫最后兩個 for 循環：

for k in range(K): # iteratively for each topic update
    n_k = np.zeros(N_W) # vocab size
    for w in range(N_W):
         a = np.logical_not(X-w) # all X(i,j) == w become a True, others a false
         b = np.logical_not(Z-k) # all Z(i,j) == w become a True, others a false
         c = np.logical_and(a,b) # all (i,j) where X(i,j) == w and Z(i,j) == k are True, others false
         n_k[w] = np.sum(c) # sum all True values

甚至作為一個班輪：

n_k = np.array([[np.sum(np.logical_and(np.logical_not(X[:N_D,:N_W]-w), np.logical_not(Z[:N_D,:N_W]-k))) for w in range(N_W)] for k in range(K)])

然后 n_k 中的每一行都可以用於 beta 計算。 現在它還包括 N_W 和 N_D 作為限制，如果它們不等於 X 和 Z 的大小

Answer 2

我對以下矩陣進行了一些測試：

import numpy as np

K = 90
N_W = 100
N_D = 11
N_W = 12

Z = np.random.randint(0, K, size=(N_D, N_W))
X = np.random.randint(0, N_W, size=(N_D, N_W))

gamma = 1

原代碼：

%%timeit
Beta = numpy.zeros((K, N_W))
for k in range(K): # iteratively for each topic update
    n_k = np.zeros(N_W) # vocab size

    for w in range(N_W):
        for i in range(N_D):
            for j in range(N_W): 
                # counting number of times a word is assigned to a topic
                n_k[w] += (X[i,j] == w) and (Z[i,j] == k) 

    # update
    Beta[k,:] = np.random.dirichlet(gamma + n_k)

865 ms ± 8.37 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

然后僅對內部兩個循環進行矢量化：

%%timeit
Beta = numpy.zeros((K, N_W))

for k in range(K): # iteratively for each topic update
    n_k = np.zeros(N_W) # vocab size

    for w in range(N_W):
        n_k[w] = np.sum((X == w) & (Z == k))


    # update
    Beta[k,:] = np.random.dirichlet(gamma + n_k)

21.6 ms ± 542 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

最后是一些廣播和提取常見元素的創造性應用：

%%timeit
Beta = numpy.zeros((K, N_W))

w = np.arange(N_W)
X_eq_w = np.equal.outer(X, w)

for k in range(K): # iteratively for each topic update
    n_k = np.sum(X_eq_w & (Z == k)[:, :, None], axis=(0, 1))


    # update
    Beta[k,:] = np.random.dirichlet(gamma + n_k)

4.6 ms ± 92.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

這里的權衡是在速度和 memory 之間。 對於我使用的形狀，這不是內存密集型，但我在最后一個解決方案中構建的中間三維 arrays 可能會變得相當大。