时间

Question

我想在numpy中做两个2d数组的元素外部产品。

A.shape = (100, 3) # A numpy ndarray
B.shape = (100, 5) # A numpy ndarray

C = element_wise_outer_product(A, B) # A function that does the trick
C.shape = (100, 3, 5) # This should be the result
C[i] = np.outer(A[i], B[i]) # This should be the result

以下是一个天真的实现。

tmp = []
for i in range(len(A):
    outer_product = np.outer(A[i], B[i])
    tmp.append(outer_product)
C = np.array(tmp)

从堆栈溢出中获得灵感的更好解决方案。

big_outer = np.multiply.outer(A, B)
tmp = np.swapaxes(tmp, 1, 2)
C_tmp = [tmp[i][i] for i in range(len(A)]
C = np.array(C_tmp)

我正在寻找一个摆脱for循环的矢量化实现。有没有人有想法？谢谢！

Answer 1

将A和B扩展为3D，保持第一个轴对齐，并分别沿第三个和第二个轴引入新轴None/np.newaxis，然后相互相乘。这将允许broadcasting为矢量化解决方案发挥作用。

因此，实现将是 -

A[:,:,None]*B[:,None,:]

我们可以使用ellipsis对A＆＃39; s :,:进行缩短，并跳过列出剩余的最后一个轴B，就像这样 -

A[...,None]*B[:,None]

作为另一种向量化方法，我们也可以使用np.einsum，一旦我们通过字符串表示法语法，并且认为这些符号是参与天真循环实现的迭代器的代表，这可能更直观，如此 - < / p>

np.einsum('ij,ik->ijk',A,B)

Answer 2

另一种使用np.lib.stride_tricks.as_strided()的解决方案。.

这里的策略本质上是构建(100, 3, 5)数组As和(100, 3, 5)数组Bs，以使这些数组的常规按元素乘积将产生期望的结果。当然，由于as_strided()，我们实际上并没有构建消耗大量内存的数组。（{as_strided()就像一张蓝图，告诉NumPy 如何您如何映射原始数组中的数据以构造As和Bs。）

def outer_prod_stride(A, B):
    """stride trick"""
    a = A.shape[-1]
    b = B.shape[-1]
    d = A.strides[-1]
    new_shape = A.shape + (b,)
    As = np.lib.stride_tricks.as_strided(A, shape=new_shape, strides=(a*d, d, 0))
    Bs = np.lib.stride_tricks.as_strided(B, shape=new_shape, strides=(b*d, 0, d))
    return As * Bs

时间

def outer_prod_broadcasting(A, B):
    """Broadcasting trick"""
    return A[...,None]*B[:,None]

def outer_prod_einsum(A, B):
    """einsum() trick"""
    return np.einsum('ij,ik->ijk',A,B)

def outer_prod_stride(A, B):
    """stride trick"""
    a = A.shape[-1]
    b = B.shape[-1]
    d = A.strides[-1]
    new_shape = A.shape + (b,)
    As = np.lib.stride_tricks.as_strided(A, shape=new_shape, strides=(a*d, d, 0))
    Bs = np.lib.stride_tricks.as_strided(B, shape=new_shape, strides=(b*d, 0, d))
    return As * Bs

%timeit op1 = outer_prod_broadcasting(A, B)
2.54 µs ± 436 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit op2 = outer_prod_einsum(A, B)
3.03 µs ± 637 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit op3 = outer_prod_stride(A, B)
16.6 µs ± 5.39 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

似乎我的大步技巧解决方案比@Divkar的解决方案都慢。 ..仍然是一个值得了解的有趣方法。

numpy elementwise外部产品

2 个答案:

时间