如何通過矢量化加速 python 中的這個 DP function

Question

所以我在這里有這個定義，

DP[i,j] = f[i,j] + min(DP[i−1, j −1], DP[i−1, j], DP[i−1, j +1])

它定義了從 NxM 矩陣頂部到矩陣底部的 go 的最小應計成本。 f 中的每個單元格代表從另一個單元格前往該單元格的價值/成本（1.2、0、10 等）。

矩陣可能很大（1500x1500，它是圖像的梯度 map ），我為我的矩陣編寫的 DP 算法每次運行大約需要一秒。 該矩陣每次執行需要運行數百次，因此總程序運行時間長達幾分鍾。 這個循環大約是我瓶頸的 99%，所以我正在嘗試使用 Python/numpys 矢量化方法優化這個循環。 我只能訪問 Numpy 和 Scipy。

注意：我幾乎沒有在 python 中編程，所以解決方案可能只是顯而易見的 idk。

第一次嘗試，只是簡單的循環，這里的時間大約是每次運行 2-2.5 秒

DP = f.copy()
for r in range(2, len(DP) - 1): # Start at row 2 since row one doesn't change
    for c in range(1, len(DP[0]) - 1):
        DP[r][c] += min(DP[r - 1, c-1:c+2])

第二次嘗試，我嘗試利用一些 numpy 矢量化函數“fromiter”來一次計算整行而不是逐列，這里的時間大約是每次運行 1-1.5 秒。 我的目標是讓這個速度至少快一個數量級，但我對如何優化它感到困惑。

DP = f.copy()
for r in range(2, len(DP) - 1):
    def foo(arr):
        idx, val = arr
        if idx == 0 or idx == len(DP[[0]) - 1:
            return np.inf
        return val + min(DP[r - 1, idx - 1], DP[r - 1, idx], DP[r - 1, idx + 1])


    DP[r, :] = np.fromiter(map(foo, enumerate(DP[r, :])))

Answer 1

正如 hpaulj 所說，由於您的問題本質上是連續的，因此很難完全矢量化，盡管這似乎是可能的（每個單元格都根據行r=2的值進行更新，不同之處在於每個單元格從第 2 行考慮的三元組數以下行）所以也許你可以找到一個聰明的方法來做到這一點！

話雖如此，一個快速和半矢量化的解決方案是使用 user42541 提出的精美索引執行滑動 windows的簡潔方法，因此我們用矢量化調用替換內部循環：

indexer = np.arange(3)[:,None] + np.arange(DP.shape[1] - 2)[None,:]
for r in range(2, DP.shape[0] - 1):
    DP[r,1:-1] += np.min(DP[r-1,indexer], axis = 0)

對於 1500x1500 的整數數組，這導致相對於您的雙循環方法（您的矢量化解決方案在我的電腦中不起作用）大約兩個數量級的加速。