Numpy,通过迭代替换广播

时间:2015-11-17 16:30:28

标签: python numpy generator broadcast

我有以下代码段

def norm(x1, x2):
    return np.sqrt(((x1 - x2)**2).sum(axis=0))

def call_norm(x1, x2):
    x1 = x1[..., :, np.newaxis]
    x2 = x2[..., np.newaxis, :]
    return norm(x1, x2)

据我了解,每个x表示N维空间中的点数组,其中N是数组最终维度的大小(因此对于3空间中的点,最终维度为3 )。它插入额外的维度并使用广播来生成这些点集的笛卡尔积,从而计算所有点对之间的距离。

x = np.array([[1, 2, 3],[1, 2, 3]])
call_norm(x, x)
array([[ 0.        ,  1.41421356,  2.82842712],
       [ 1.41421356,  0.        ,  1.41421356],
       [ 2.82842712,  1.41421356,  0.        ]])

(所以[1,1][2,2]之间的距离为1.41421356,符合预期)

我发现对于中等大小的问题,这种方法可能会占用大量内存。我可以轻松地去除矢量"问题并通过迭代替换它,但我希望它会变慢。我有一个(合理的)简单的折衷解决方案,我可以拥有矢量化的大部分速度优势但没有内存损失?一些花哨的发电机技巧?

1 个答案:

答案 0 :(得分:1)

如果没有numpy矢量化的内存惩罚,就没有办法进行这种计算。对于有效计算成对距离矩阵的特定情况,包通过在C(例如scipy.spatial.distance)或Cython(例如sklearn.metrics.pairwise)中实现事物来解决这个问题。

如果你想“手动”这样做,可以说,使用numpy风格的语法,但不会导致内存损失,目前最好的选项可能是dask.array,它可以自动构建和执行使用numpy风格语法进行批量执行的灵活任务图。

以下是使用dask进行此计算的示例:

import dask.array as da

# Create the chunked data. This can be created
# from numpy arrays as well, e.g. x_dask = da.array(x_numpy)
x = da.random.random((100, 3), chunks=5)
y = da.random.random((200, 3), chunks=5)

# Compute the task graph (syntax just like numpy!)
diffs = x[:, None, :] - y[None, :, :]
dist = da.sqrt((diffs ** 2).sum(-1))

# Execute the task graph
result = dist.compute()
print(result.shape)
# (100, 200)

你会发现dask比NumPy具有更高的内存效率,通常比NumPy更具计算效率,并且也可以相对直接地在并行/外核中计算。

相关问题