基于唯一值对 2D Numpy/CuPy 数组进行更快的迭代

Posted 2023-02-16

【中文标题】基于唯一值对 2D Numpy/CuPy 数组进行更快的迭代

【英文标题】：Faster iteration over 2D Numpy/CuPy arrays based on unique values

【发布时间】：2021-04-17 18:13:10

【问题描述】：

我目前正在循环一个 numpy 数组以对其进行切片并执行一些 ndarray 数组。由于 2001*2001 元素数组的大小，目前需要的时间非常长。因此，我希望有人可能会提示，如何加速代码：

import cupy as cp
from time import time

height, width = 187, 746
org_sized = cp.random.rand(2001, 2001) * 60

height_mat = cp.random.rand(height, width) * 100 # orinally values getting larger from (0, width//2) to the outside with the distance squared

indices = cp.indices((height, width))
y_offsets = indices[0]
x_offsets = indices[1] - (width + 1)/2
angle_mat = cp.round_(2*(90 - cp.rad2deg(cp.arctan2(y_offsets, x_offsets))) + 180).astype(int)

weights = cp.random.rand(361)/ 10  # weights oroiginally larger in the middle

# pad the org_sized matrix with zeros to a fit a size of (2001+heigth, 2001+weight)
west = cp.zeros((org_sized.shape[0], width // 2))
east = cp.zeros((org_sized.shape[0], round(width // 2)))

enlarged_size = cp.hstack((west, org_sized))
enlarged_size = cp.hstack((enlarged_size, east))

south = cp.zeros((height, enlarged_size.shape[1]))

enlarged_size = cp.vstack((enlarged_size, south))

shadow_time_hrs = cp.zeros_like(org_sized)


for y in range(org_sized.shape[0]):
    start_time = time()
    for x in range(org_sized.shape[1]):
        # shift h_extras and angles that they match in size, and are correctly aligned
        short_elevations = enlarged_size[y:y+height, x:x+width]

        overshadowed = (short_elevations - org_sized[y, x]) > height_mat
        shadowed_angles = angle_mat * overshadowed
        shadowed_segments = cp.unique(shadowed_angles)
        angle_segments = shadowed_segments

        sum_hours = cp.sum(weights[angle_segments])
        shadow_time_hrs[y, x] = sum_hours
    if (y % 100) == 0:
        print(f"Computation for line y took: time() - start_time.")

首先我在函数 calc_shadow_point 上使用了 numbas @njit，但结果证明它比不使用时慢 2 倍。因此我将numpy数组切换到cupy数组。这提供了大约 50% 的加速。可能是因为数组太小了。

对于此类问题，除了迭代之外，还有其他方法吗，或者有没有一种方法可以在迭代器上使用多线程进行迭代？

编辑：我将代码更改为相同运行时的最小示例（每行 org_sized 1.1 秒）。不知何故，我必须提高计算速度。低于当前计算时间 10% 的所有内容都将使代码可用。 由于评论，我将 np.unique 更改为 cp.unique，但正如所言。它并没有导致仅 6 % 的大幅加速。我目前正在使用 GTX 1060。但是当它有帮助时可以设法使用 1660 Ti。

【问题讨论】：

我怀疑深度嵌入的unique 是瓶颈。这对其输入进行排序并检查重复项。

“到 numpy 数组到 cupy 数组”你的意思是 from numpy 数组 to cupy 数组？请注意，在 GPU 上有效地实现np.unique 非常困难，因此它的效率不应该很高。 width 和 height 的值是多少？更重要的是shadowed_angles 中的值是什么？请注意，拥有MWE 应该会对我们有所帮助。

想法是：拥有一个更大的数组（称为海拔），大小为 (1, 2001,20001)。并在其上使用一些“过滤器”。过滤器具有（高度，宽度）的形状，该形状在启动时确定，而不是固定。在我当前的示例中，它是 (187, 746)。但是这个比率是固定的（除了四舍五入）。 showded_angles 的当前类型是 int32。但是因为它包含 0 到 360 的值（所有整数），我也可以将其设置为“uint16”，但是。我想，这不会让想法更好吗？矩阵angle存储了角段的信息。

【参考方案1】：

unique 很慢（在 CPU 和 GPU 上），因为它通常在内部使用哈希映射或排序。此外，正如您所说，阵列太小而无法在 GPU 上高效，从而导致巨大的内核开销。希望您不需要它：您可以使用 bincount（带有minlength=361 和一个扁平数组），因为您知道这些值是有界中的小正整数范围0:361。实际上，您实际上并不需要像bincount 那样计算这些值，您只想知道0:361 范围内的哪些值存在于shadowed_angles 中。因此，可以使用 Numba 编写更快的bincount 实现。此外，数组计算可以连续完成减少了分配量和内存压力。最后，parallelism可以用来加速计算（使用Numba的prange和parallel=True）。

这是生成的基于CPU的实现：

@nb.njit
def computeSumHours(org_sized, enlarged_size, angle_mat, height_mat, shadow_time_hrs, y, x):
    height, width = height_mat.shape
    short_elevations = enlarged_size[y:y+height, x:x+width]
    shadowed_segments = np.zeros(361)

    for y2 in range(height):
        for x2 in range(width):
            overshadowed = (short_elevations[y2, x2] - org_sized[y, x]) > height_mat[y2, x2]
            shadowed_angle = angle_mat[y2, x2] * overshadowed
            shadowed_segments[shadowed_angle] = weights[shadowed_angle]

    return shadowed_segments.sum()

@nb.njit(parallel=True)
def computeLine(org_sized, enlarged_size, angle_mat, height_mat, shadow_time_hrs, y):
    height, width = height_mat.shape

    for x in nb.prange(org_sized.shape[1]):
        shadow_time_hrs[y, x] = computeSumHours(org_sized, enlarged_size, angle_mat, height_mat, shadow_time_hrs, y, x)

def computeAllLines(org_sized, enlarged_size, angle_mat, height_mat, shadow_time_hrs):
    height, width = height_mat.shape

    for y in range(org_sized.shape[0]):
        start_time = time()
        computeLine(org_sized, enlarged_size, angle_mat, height_mat, shadow_time_hrs, y)
        if (y % 100) == 0:
            print("Computation for line %d took: %f." % (y, time() - start_time))

computeAllLines(org_sized, enlarged_size, angle_mat, height_mat, shadow_time_hrs)

这是我机器上每次迭代的计时结果（使用 i7-9600K 和 GTX-1660-Super）：

Reference implementation (CPU): 2.015 s
Reference implementation (GPU): 0.882 s
Optimized implementation (CPU): 0.082 s

这比基于 GPU 的参考实现快 10 倍，比基于 CPU 的参考实现快 25 倍。

请注意，可以在 GPU 上使用相同的技术，但不能使用 CuPy：需要编写一个 GPU 内核来执行此操作（例如，使用 CUDA）。但是，要有效地做到这一点是相当复杂的。

【讨论】：

非常感谢。它在我的机器上也快 10 倍，比我自己尝试使用 numba 的速度快 50 倍！太棒了。

@OrdensRitter 欢迎您。不要忘记验证答案；）。