简化循环 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 重塑和转换参数，将任务简化为 dot。因此，虽然可以切换哪个轴获得乘积和，但它处理 batches 的效果并不比 dot 好多少。这是添加 matmul 的重要原因。 np.einsum 提供相同的功能（甚至更多），但性能通常不是那么好（除非它被“优化”为等效的 matmul）。

简化循环 np.tensordot 表达式

Simplifying looped np.tensordot expression

python

numpy

tensordot