时间戳

首页 > 解决方案 > 基于唯一值对 2D Numpy/CuPy 数组进行更快的迭代

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

问题描述

我目前正在循环一个 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% 的加速。Probapy，因为数组太小了。

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

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

标签： pythonperformancenumpycupy

解决方案

unique很慢（在 CPU 和 GPU 上），因为它通常在内部使用哈希映射或排序。此外，正如您所说，阵列太小而无法在 GPU 上高效，从而导致巨大的内核开销。希望您不需要它：您可以使用bincount(with minlength=361and a flatten array)，因为您知道这些值是有界范围内的小正整数0:361。实际上，您实际上并不需要像 do 那样计算值bincount，您只想知道范围的哪些值0:361存在于shadowed_angles. 因此，bincount可以使用Numba编写更快的实现。此外，数组计算可以在一行中完成，从而减少了分配和内存压力。最后，并行性可用于加速计算（使用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）。然而，要有效地做到这一点是相当复杂的。

推荐阅读