广播是个坏主意吗？(numpy的)

Question

术语广播描述了numpy如何在算术运算期间处理具有不同形状的数组.

Example 1:

from numpy import array 
a = array([1.0,2.0,3.0])
b = array([2.0,2.0,2.0]) # multiply element-by-element ()
a * b
>> array([ 2.,  4.,  6.])

Example 2 :

from numpy import array
a = array([1.0,2.0,3.0])
b = 2.0  # broadcast b to all a
a * b
>>array([ 2.,  4.,  6.])

我们可以将在算术运算期间被拉伸的标量b想象成具有与a相同形状的数组.Numpy足够聪明,可以使用原始标量值,而无需实际制作副本,以便广播操作尽可能具有内存和计算效率(b是标量,而不是数组)

@Eric Duminil在另一个内存性能问题中做出的一个小基准测试表明,广播在速度和内存方面存在差异

但是,我引用上面链接的同一篇文章:

有些情况下,广播是一个坏主意,因为它会导致内存的低效使用,从而减慢计算速度

问题是:当广播使用不必要的大量内存并导致性能低下时？换句话说,当我们应该使用混合广播/python循环算法而不是纯粹的广播approch？

Answer 1

广播质量不佳的情况并没有真正明确的案例。通常，广播是最简单、最具可读性的解决方案，这可能就是您想要的。如果在基准测试后速度太慢，我只会考虑优化广播或顺序操作。

正如许多现有评论所说，在广播方面，内存和计算之间通常需要权衡。然而，如果你的算法设计不正确，你可能会损害这两个方面。

• 我发现的最大问题是，虽然 numpy 可能会尝试优化不同的步骤，以便使用数组的视图，但通常它无法对广播或顺序操作进行这些优化。巧妙使用 numpy 可能仍然无法解决这个问题，因此可能值得考虑使用循环重写程序，以便您可以手动将操作合并在一起。这可以最大限度地减少内存使用并最大限度地提高性能。然而，在普通 python 中执行此操作会非常慢，但幸运的是，我们有类似的东西numba可以 JIT（及时）将带注释的函数编译成高效的机器代码。

• 另一个问题是数组非常大，广播会迅速增加内存使用量。如果数组的形状不同，通常会从 O(n) 甚至 O(n^2) 移动，甚至更糟（我不是说用标量进行广播）。虽然这对于小型阵列来说可能没问题，但随着阵列大小的增加，它很快就会成为一个问题。乘以定标器可能只会使内存使用量增加一倍，但这并没有那么糟糕。

  a = np.arange(128, dtype='float64')
  
  # temp. array memory usage -- A: ~1KB, B: ~1MB, C: ~1GB
  A: float = a.sum() 
  B: float = (a[:, None] + a[None, :]).sum()
  C: float = (a[:, None, None] + a[None, :, None] + a[None, None, :]).sum()
  
  # If the size of a was small at, then we are OK
  # eg. a.shape == (16,) gives -- a: ~128B, b: ~2KB, c: ~32KB
  

• 上面的示例虽然可读性不高，但由于归约操作中使用临时数组而导致数组大小增加，因此可以将其转换为仅使用 O(n) 的基于循环的格式。如果您想要的输出是广播数组本身而不是减少，那么广播将非常接近最佳。

示例：我最近自己也遇到了这个问题。我有一个一维二进制掩码，我需要将其自身广播到二维矩阵中，以便我可以使用它从大型预先计算的距离矩阵中提取元素，额外的条件是我也必须排除对角线（我也无法访问原始的 1d 位置）。

自然，这看起来如下：

import numpy as np

def broadcast_masked_tril_total(dists2d, mask):
    # broadcast 1d mask into 2d array
    # - this can be very slow, moving from O(N) to O(N^2) memory
    mask2d = mask[None, :] & mask[:, None]
    # ignore diagonal
    # - the 2D array needs to exist in memory to make these edits, a view cannot work.
    np.fill_diagonal(mask2d, False)
    # index array with elements
    # - 2d mask means O(N^2) memory is read, instead of O(N)
    total = dists2d[mask2d].sum()
    # elems are repeated so divide by 2
    return total / 2

问题当然是内存使用和值的中间存储。

可能有一个聪明的 numpy 修复，但明显的解决方案只是将其转换为循环，正如您所说，您不需要使用广播。有利的是，您可以尝试确定哪些操作可以合并在一起，而不是像 numpy 那样将它们链接起来。

通常的一般经验法则是，内存访问和中间存储位置越少，运行速度就越快。

import numba

@numba.njit()
def efficient_masked_tril_total(dists2d, mask):
    total = 0.
    for y, y_val in enumerate(mask):
        # we only ever need to read from the same O(N) mask memory
        # we can immediately skip invalid rows
        if not y_val:
            continue
        # skip the diagonal and one triangle of the distance matrix
        # - can't do this efficiently with numpy broadcasting and
        #   mask, intermediate storage of 2d mask was required
        for x in range(y+1, len(mask)):
            # again accessing the same O(n) mask item without broadcasting
            if not mask[x]:
                continue
            total += dists2d[y, x]
    return total

例如使用这个：

N = int(np.sqrt((10*1024**2)/(64/8)))  # enough elems for 10.0 MB
# make distance matrices
mask = np.random.random(N) > 0.5
positions = np.random.random(N)
# again we broadcast, note that we could further optimise
# our efficient approach by just passing in the positions directly.
dists2d = np.abs(positions[:, None] - positions[None, :])
# warmup
print(broadcast_masked_tril_total(dists2d, mask))
print(efficient_masked_tril_total(dists2d, mask))
# timeit
import timeit
print(timeit.timeit(lambda: broadcast_masked_tril_total(dists2d, mask), number=1000))
print(timeit.timeit(lambda: efficient_masked_tril_total(dists2d, mask), number=1000))

tl;dr：简而言之，我建议您始终使用最简单、最具可读性的解决方案。只有当它成为性能问题时，您才应该花时间进行基准测试和优化您的方法。请记住“过早的优化是万恶之源”。

因此，在具体情况下，广播并不是一个坏主意。通常这是最简单的解决方案，这是一件好事。有时，如果您需要返回实际的广播数组，则广播将是最佳解决方案。如果您将广播与某种辅助操作（例如减少）一起使用，则可以通过将其转换为循环来优化组合操作。请记住，这不一定是坏事，只有当性能成为问题时才成为问题。较小的数组通常不是问题，但如果您使用较大的数组，广播也很容易导致内存问题。