python - 简化循环的 np.tensordot 表达式

Question

目前，我的脚本如下所示：

import numpy as np

a = np.random.rand(10,5,2)
b = np.random.rand(10,5,50)
c = np.random.rand(10,2,50)

for i in range(a.shape[0]):
    c[i] = np.tensordot(a[i], b[i], axes=(0,0))

我想在不使用 for 循环的情况下复制相同的行为，因为它可以并行完成。但是，我还没有找到使用 tensordot 函数执行此操作的简洁方法。有没有办法为这个操作创建一个单线？

Answer 1

在这种情况下，您可以使用 numpy.einsum 函数

c = np.einsum('ijk,ijl->ikl', a, b)

Answer 2

的替代方案einsum是matmul/@。第一个数组必须被转置，所以乘积总和维度是最后一个：

In [162]: a = np.random.rand(10,5,2)
     ...: b = np.random.rand(10,5,50)
In [163]: c=a.transpose(0,2,1)@b
In [164]: c.shape
Out[164]: (10, 2, 50)
In [165]: c1 = np.random.rand(10,2,50)
     ...: 
     ...: for i in range(a.shape[0]):
     ...:     c1[i] = np.tensordot(a[i], b[i], axes=(0,0))
     ...: 
In [166]: np.allclose(c,c1)
Out[166]: True

tensordot重塑和转置参数，将任务简化为 simple dot。因此，虽然可以切换哪些轴获得乘积和，但它的处理能力并不batches比dot. matmul这是添加 原因的很大一部分。np.einsum提供相同的功率（甚至更多），但性能通常不是那么好（除非它已“优化”到等效的matmul）。