numpy - numpy elementwise 外积

Question

我想在numpy中做两个二维数组的元素外积。

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我们可以通过使用for A来缩短它::,:并跳过用 列出剩余的最后一个轴B，就像这样 -

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

---

作为另一种向量化的方法，我们也可以使用np.einsum，一旦我们超越了字符串符号语法并认为这些符号代表了幼稚循环实现中涉及的迭代器，这可能会更直观，就像这样 -

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 的解决方案都慢。..仍然是一个值得了解的有趣方法。