简化循环的 np.tensordot 表达式答案

【问题标题】：Simplifying looped np.tensordot expression简化循环的 np.tensordot 表达式
【发布时间】：2021-11-20 12:31:00
【问题描述】：

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

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 函数执行此操作的简洁方法。有没有办法为这个操作创建一个单行？

【问题讨论】：

标签： python numpy tensordot

【解决方案1】：

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

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

【讨论】：

【解决方案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）。

【讨论】：