python - 替换 Python 中的深度循环

Question

不幸的是，我的 Python 代码中经常有 while 循环，这会导致我的程序显着变慢。

下面是一个 while 循环的例子(shape = (1000,1000,3)):

i = 0
j = 0
while i < arr.shape[0]:
    while j < arr.shape[1]:
        if arr[i,j,0] <= 5 and arr[i,j,0] > 0:
            arr[i,j,:] = 1
        else:
            arr[i,j,:] = 0

        j = j + 1
    j = 0
    i = i + 1

-->

def f(x):
        return 1 if x <= 5and x > 0 else 0

f = np.vectorize(f)
arr= f(arr)

编辑:另一个 while 循环

i = 0
j = 0
while i < arr1.shape[0]:
    while j < arr1.shape[1]:

        if arr1[i, j, 0] == 0 and arr1[i, j, 1] == 0 and arr1[i, j, 2] == 0:
            arr1[i, j, :] = arr1[i, j, :]
        else:
            arr1[i, j, :] = arr2[i, j, :]

        j = j + 1
    j = 0
    i = i + 1

有没有办法加快速度？我不确定如何。

Answer 1

编辑匆忙中，我之前的回答是错误的。感谢@gboffi 指出。

第 1 部分

def original(arr):
    i = 0
    j = 0
    while i < arr.shape[0]:
        while j < arr.shape[1]:
            if arr[i,j,0] <= 5 and arr[i,j,0] > 0:
                arr[i,j,:] = 1
            else:
                arr[i,j,:] = 0

            j = j + 1
        j = 0
        i = i + 1

    return arr

def vectorized(arr):
    mask = (0 < arr[:, :, 0]) & (arr[:, :, 0] <= 5)
    arr[mask] = 1
    arr[~mask] = 0
    return arr

def vectorized2(arr):
    """Works only if assigning 0 and 1s"""
    mask = (0 < arr[:, :, 0]) & (arr[:, :, 0] <= 5)
    mask = np.dstack([mask] * arr.shape[2])
    return mask.astype(np.float32)

基准测试

以下使用下一个基准测试中的 arr。我没有时间对此进行广泛的测试，所以我建议您运行下面的矢量化版本并使用 np.allclose 比较原始代码在 arr< 的各种情况下的结果 您可以验证。

In [101]: %timeit v1 = vectorized(arr.copy())
1000 loops, best of 3: 312 µs per loop

In [102]: %timeit v2 = original(arr.copy())
100 loops, best of 3: 10.4 ms per loop

In [103]: np.allclose(v1, v2)
Out[103]: True

In [108]: %timeit v3 = vectorized2(arr.copy())
10000 loops, best of 3: **83.3 µs** per loop

In [110]: v3 = vectorized2(arr.copy())

In [111]: np.allclose(v1, v3)
Out[111]: True

第 2 部分

def vectorized(arr1, arr2):
    mask = np.all(arr1 == 0, axis=2)
    mask = np.dstack([mask] * arr1.shape[2])
    return np.where(mask, arr1, arr2)

def original(arr1, arr2):
    i = 0
    j = 0
    while i < arr1.shape[0]:
        while j < arr1.shape[1]:

            if arr1[i, j, 0] == 0 and arr1[i, j, 1] == 0 and arr1[i, j, 2] == 0:
                arr1[i, j, :] = arr1[i, j, :]
            else:
                arr1[i, j, :] = arr2[i, j, :]

            j = j + 1

        j = 0
        i = i + 1

    return arr1

第二个循环的基准

# Prepare data
m = 100
n = 100
d = 3
np.random.seed(0)
arr = np.random.randint(0, 11, size=(m, n, d))

true_mask = np.random.randint(0, 2, size=(m, n, 1), dtype=np.bool)
true_mask = np.dstack([true_mask] * d)

arr[true_mask] = 0

arr1 = arr.copy()
arr2 = -1 * np.ones_like(arr1)

In [84]: v1 = vectorized(arr1, arr2)

In [85]: v2 = original(arr1, arr2)

In [86]: np.allclose(v1, v2)
Out[86]: True

In [87]: %timeit v1 = vectorized(arr1, arr2)
1000 loops, best of 3: 284 µs per loop

In [88]: %timeit v2 = original(arr1, arr2)
100 loops, best of 3: 12.6 ms per loop