问题:了解NumPy的einsum
我正在努力了解确切的einsum
工作原理。我看了一下文档和一些示例,但看起来似乎并不固定。
这是我们在课堂上讲的一个例子:
C = np.einsum("ij,jk->ki", A, B)
对于两个数组A
和B
我认为可以A^T * B
,但是我不确定(正在对其中之一进行移调吗?)。谁能告诉我这里到底发生了什么(以及使用时的一般情况einsum
)?
回答 0
(注:这个答案是基于短的博客文章约einsum
我写了前一阵子。)
怎么einsum
办?
假设我们有两个多维数组,A
和B
。现在假设我们要…
- 乘
A
用B
一种特殊的方式来创造新的产品阵列; 然后也许 - 沿特定轴求和该新数组;然后也许
- 以特定顺序转置新数组的轴。
有一个很好的机会,einsum
可以帮助我们做到这一点更快,内存更是有效的NumPy的功能组合,喜欢multiply
,sum
和transpose
允许。
einsum
工作如何?
这是一个简单(但并非完全无关紧要)的示例。取以下两个数组:
A = np.array([0, 1, 2])
B = np.array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
我们将逐个元素相乘A
,B
然后沿着新数组的行求和。在“普通” NumPy中,我们将编写:
>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])
因此,此处的索引操作A
将两个数组的第一个轴对齐,以便可以广播乘法。然后将乘积数组中的行相加以返回答案。
现在,如果我们想使用它einsum
,我们可以这样写:
>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])
该签名字符串'i,ij->i'
是这里的关键,需要解释的一点。您可以将其分为两半。在左侧(的左侧->
),我们标记了两个输入数组。在的右侧->
,我们标记了要结束的数组。
接下来会发生以下情况:
A
有一个轴 我们已经标记了它i
。并且B
有两个轴;我们将轴0标记为i
,将轴1 标记为j
。通过在两个输入数组中重复标签
i
,我们告诉我们einsum
这两个轴应该相乘。换句话说,就像A
数组B
一样,我们将array 与array 的每一列相乘A[:, np.newaxis] * B
。请注意,
j
它不会在所需的输出中显示为标签;我们刚刚使用过i
(我们想以一维数组结尾)。通过省略标签,我们告诉einsum
来总结沿着这条轴线。换句话说,我们就像对行进行求和.sum(axis=1)
。
基本上,这是您需要了解的所有信息einsum
。玩一会会有所帮助;如果我们将两个标签都留在输出中,则会'i,ij->ij'
返回2D产品数组(与相同A[:, np.newaxis] * B
)。如果我们说没有输出标签,'i,ij->
我们将返回一个数字(与相同(A[:, np.newaxis] * B).sum()
)。
einsum
但是,最重要的是,它不会首先构建临时产品系列;它只是对产品进行累加。这样可以节省大量内存。
一个更大的例子
为了解释点积,这里有两个新数组:
A = array([[1, 1, 1],
[2, 2, 2],
[5, 5, 5]])
B = array([[0, 1, 0],
[1, 1, 0],
[1, 1, 1]])
我们将使用计算点积np.einsum('ij,jk->ik', A, B)
。这是一张图片,显示了从函数获得的A
和B
和输出数组的标签:
您会看到j
重复的标签-这意味着我们会将的行A
与的列相乘B
。此外,j
输出中不包含标签-我们对这些产品进行求和。标签i
和k
被保留用于输出,因此我们得到一个2D数组。
这一结果与其中标签阵列比较可能是更加明显j
的不求和。在下面的左侧,您可以看到写入产生的3D数组np.einsum('ij,jk->ijk', A, B)
(即,我们保留了label j
):
求和轴j
给出了预期的点积,如右图所示。
一些练习
为了获得更多的感觉einsum
,使用下标符号实现熟悉的NumPy数组操作可能会很有用。任何涉及乘法和求和轴组合的内容都可以使用编写 einsum
。
令A和B为两个具有相同长度的一维数组。例如A = np.arange(10)
和B = np.arange(5, 15)
。
的总和
A
可以写成:np.einsum('i->', A)
A * B
可以按元素写成:np.einsum('i,i->i', A, B)
内积或点积
np.inner(A, B)
或np.dot(A, B)
可以写成:np.einsum('i,i->', A, B) # or just use 'i,i'
外部乘积
np.outer(A, B)
可以写成:np.einsum('i,j->ij', A, B)
对于2D数组,C
和D
,只要轴是兼容的长度(相同长度或其中之一具有长度1),下面是一些示例:
C
(主对角线总和)的轨迹np.trace(C)
可以写成:np.einsum('ii', C)
的元素方式乘法
C
和转置D
,C * D.T
可以写成:np.einsum('ij,ji->ij', C, D)
可以将每个元素乘以
C
该数组D
(以构成4D数组)C[:, :, None, None] * D
,可以写成:np.einsum('ij,kl->ijkl', C, D)
回答 1
numpy.einsum()
如果您直观地理解它的想法,将非常容易。作为示例,让我们从涉及矩阵乘法的简单描述开始。
使用时numpy.einsum()
,您要做的就是传递所谓的下标字符串作为参数,然后传递输入数组。
假设您有两个2D数组A
和B
,并且想要进行矩阵乘法。所以你也是:
np.einsum("ij, jk -> ik", A, B)
在这里,下标字符串 ein
此之后的下标字符串->
将成为我们的结果数组。如果将其保留为空,则将对所有内容求和,并返回标量值作为结果。否则,所得数组将具有根据下标字符串的尺寸。在我们的示例中,它将为ik
。这很直观,因为我们知道对于矩阵乘法,数组中的列数A
必须与数组中的行数相匹配,B
这就是这里发生的情况(即,我们通过在下标字符串中重复char j
来编码此知识)
这里还有一些其他示例,简要说明了np.einsum()
实现某些常见张量或nd数组操作的用途/功能。
输入项
# a vector
In [197]: vec
Out[197]: array([0, 1, 2, 3])
# an array
In [198]: A
Out[198]:
array([[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]])
# another array
In [199]: B
Out[199]:
array([[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4]])
1)矩阵乘法(类似于np.matmul(arr1, arr2)
)
In [200]: np.einsum("ij, jk -> ik", A, B)
Out[200]:
array([[130, 130, 130, 130],
[230, 230, 230, 230],
[330, 330, 330, 330],
[430, 430, 430, 430]])
2)沿主对角线提取元素(类似于np.diag(arr)
)
In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])
3)Hadamard乘积(即两个数组的按元素乘积)(类似于arr1 * arr2
)
In [203]: np.einsum("ij, ij -> ij", A, B)
Out[203]:
array([[ 11, 12, 13, 14],
[ 42, 44, 46, 48],
[ 93, 96, 99, 102],
[164, 168, 172, 176]])
4)逐元素平方(类似于np.square(arr)
或arr ** 2
)
In [210]: np.einsum("ij, ij -> ij", B, B)
Out[210]:
array([[ 1, 1, 1, 1],
[ 4, 4, 4, 4],
[ 9, 9, 9, 9],
[16, 16, 16, 16]])
5)痕迹(即主对角元素的总和)(类似于np.trace(arr)
)
In [217]: np.einsum("ii -> ", A)
Out[217]: 110
6)矩阵转置(类似于np.transpose(arr)
)
In [221]: np.einsum("ij -> ji", A)
Out[221]:
array([[11, 21, 31, 41],
[12, 22, 32, 42],
[13, 23, 33, 43],
[14, 24, 34, 44]])
7)(向量的)外积(类似于np.outer(vec1, vec2)
)
In [255]: np.einsum("i, j -> ij", vec, vec)
Out[255]:
array([[0, 0, 0, 0],
[0, 1, 2, 3],
[0, 2, 4, 6],
[0, 3, 6, 9]])
8)(向量的)内积(类似于np.inner(vec1, vec2)
)
In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14
9)沿轴0求和(类似于np.sum(arr, axis=0)
)
In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])
10)沿轴1的总和(类似于np.sum(arr, axis=1)
)
In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4, 8, 12, 16])
11)批矩阵乘法
In [287]: BM = np.stack((A, B), axis=0)
In [288]: BM
Out[288]:
array([[[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]],
[[ 1, 1, 1, 1],
[ 2, 2, 2, 2],
[ 3, 3, 3, 3],
[ 4, 4, 4, 4]]])
In [289]: BM.shape
Out[289]: (2, 4, 4)
# batch matrix multiply using einsum
In [292]: BMM = np.einsum("bij, bjk -> bik", BM, BM)
In [293]: BMM
Out[293]:
array([[[1350, 1400, 1450, 1500],
[2390, 2480, 2570, 2660],
[3430, 3560, 3690, 3820],
[4470, 4640, 4810, 4980]],
[[ 10, 10, 10, 10],
[ 20, 20, 20, 20],
[ 30, 30, 30, 30],
[ 40, 40, 40, 40]]])
In [294]: BMM.shape
Out[294]: (2, 4, 4)
12)沿轴2的总和(类似于np.sum(arr, axis=2)
)
In [330]: np.einsum("ijk -> ij", BM)
Out[330]:
array([[ 50, 90, 130, 170],
[ 4, 8, 12, 16]])
13)对数组中的所有元素求和(类似于np.sum(arr)
)
In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480
14)多轴总和(即边际化)
(类似于np.sum(arr, axis=(axis0, axis1, axis2, axis3, axis4, axis6, axis7))
)
# 8D array
In [354]: R = np.random.standard_normal((3,5,4,6,8,2,7,9))
# marginalize out axis 5 (i.e. "n" here)
In [363]: esum = np.einsum("ijklmnop -> n", R)
# marginalize out axis 5 (i.e. sum over rest of the axes)
In [364]: nsum = np.sum(R, axis=(0,1,2,3,4,6,7))
In [365]: np.allclose(esum, nsum)
Out[365]: True
15)双点积(类似于np.sum(哈达玛积) cf. 3)
In [772]: A
Out[772]:
array([[1, 2, 3],
[4, 2, 2],
[2, 3, 4]])
In [773]: B
Out[773]:
array([[1, 4, 7],
[2, 5, 8],
[3, 6, 9]])
In [774]: np.einsum("ij, ij -> ", A, B)
Out[774]: 124
16)2D和3D阵列乘法
在要验证结果的线性方程组(Ax = b)求解时,这种乘法可能非常有用。
# inputs
In [115]: A = np.random.rand(3,3)
In [116]: b = np.random.rand(3, 4, 5)
# solve for x
In [117]: x = np.linalg.solve(A, b.reshape(b.shape[0], -1)).reshape(b.shape)
# 2D and 3D array multiplication :)
In [118]: Ax = np.einsum('ij, jkl', A, x)
# indeed the same!
In [119]: np.allclose(Ax, b)
Out[119]: True
相反,如果必须使用reshape
操作才能获得相同的结果,例如:
# reshape 3D array `x` to 2D, perform matmul
# then reshape the resultant array to 3D
In [123]: Ax_matmul = np.matmul(A, x.reshape(x.shape[0], -1)).reshape(x.shape)
# indeed correct!
In [124]: np.allclose(Ax, Ax_matmul)
Out[124]: True
回答 2
让我们制作2个数组,它们具有不同但兼容的维度,以突出它们之间的相互作用
In [43]: A=np.arange(6).reshape(2,3)
Out[43]:
array([[0, 1, 2],
[3, 4, 5]])
In [44]: B=np.arange(12).reshape(3,4)
Out[44]:
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
您的计算将(2,3)的“点”(乘积之和)与(3,4)相乘,以生成(4,2)数组。 i
是第一个昏暗的A
,最后一个C
; k
最后B
1个,第1个C
。 j
通过求和“消耗”。
In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]:
array([[20, 56],
[23, 68],
[26, 80],
[29, 92]])
这与np.dot(A,B).T
-是转置的最终输出相同。
要查看更多情况j
,请将C
下标更改为ijk
:
In [46]: np.einsum('ij,jk->ijk',A,B)
Out[46]:
array([[[ 0, 0, 0, 0],
[ 4, 5, 6, 7],
[16, 18, 20, 22]],
[[ 0, 3, 6, 9],
[16, 20, 24, 28],
[40, 45, 50, 55]]])
这也可以通过以下方式生成:
A[:,:,None]*B[None,:,:]
即,添加一个k
维度的端部A
,以及i
与前部B
,产生了(2,3,4)阵列。
0 + 4 + 16 = 20
,9 + 28 + 55 = 92
等; 求和j
转置以获得较早的结果:
np.sum(A[:,:,None] * B[None,:,:], axis=1).T
# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,) j,k]
回答 3
我发现NumPy:交易技巧(第二部分)具有启发性
我们使用->指示输出数组的顺序。因此,将“ ij,i-> j”视为具有左侧(LHS)和右侧(RHS)。LHS上标签的任何重复都会明智地计算乘积元素,然后求和。通过更改RHS(输出)端的标签,我们可以相对于输入数组定义要在其中进行处理的轴,即沿轴0、1求和,依此类推。
import numpy as np
>>> a
array([[1, 1, 1],
[2, 2, 2],
[3, 3, 3]])
>>> b
array([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
>>> d = np.einsum('ij, jk->ki', a, b)
请注意,存在三个轴,即i,j,k,并且重复了j(在左侧)。 i,j
代表的行和列a
。j,k
为b
。
为了计算乘积并对齐j
轴,我们需要在上添加一个轴a
。(b
将沿第一个轴广播?)
a[i, j, k]
b[j, k]
>>> c = a[:,:,np.newaxis] * b
>>> c
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8]],
[[ 0, 2, 4],
[ 6, 8, 10],
[12, 14, 16]],
[[ 0, 3, 6],
[ 9, 12, 15],
[18, 21, 24]]])
j
在右侧不存在,因此我们求和j
是3x3x3数组的第二个轴
>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
[18, 24, 30],
[27, 36, 45]])
最后,索引在右侧(按字母顺序)相反,因此我们进行了转置。
>>> c.T
array([[ 9, 18, 27],
[12, 24, 36],
[15, 30, 45]])
>>> np.einsum('ij, jk->ki', a, b)
array([[ 9, 18, 27],
[12, 24, 36],
[15, 30, 45]])
>>>
回答 4
在阅读einsum方程式时,我发现最简单的方法就是将它们简化为必要的形式。
让我们从以下(强加)语句开始:
C = np.einsum('bhwi,bhwj->bij', A, B)
首先通过标点符号进行操作,我们看到在箭头前有两个4个逗号分隔的斑点- bhwi
和bhwj
,在箭头后有一个3个字母斑点bij
。因此,该方程从两个4级张量输入产生3级张量结果。
现在,让每个斑点中的每个字母成为范围变量的名称。字母在Blob中出现的位置是该张量范围内的轴的索引。因此,产生C的每个元素的命令式求和必须从三个嵌套的for循环开始,每个C的索引一个。
for b in range(...):
for i in range(...):
for j in range(...):
# the variables b, i and j index C in the order of their appearance in the equation
C[b, i, j] = ...
因此,从本质for
上讲,每个C的输出索引都有一个循环。我们现在暂时不确定范围。
接下来,我们看一下左侧-是否有没有出现在右侧的范围变量?在我们的情况下-是,h
并且w
。for
为每个此类变量添加一个内部嵌套循环:
for b in range(...):
for i in range(...):
for j in range(...):
C[b, i, j] = 0
for h in range(...):
for w in range(...):
...
现在,在最内层的循环中,我们定义了所有索引,因此我们可以编写实际的求和并完成转换:
# three nested for-loops that index the elements of C
for b in range(...):
for i in range(...):
for j in range(...):
# prepare to sum
C[b, i, j] = 0
# two nested for-loops for the two indexes that don't appear on the right-hand side
for h in range(...):
for w in range(...):
# Sum! Compare the statement below with the original einsum formula
# 'bhwi,bhwj->bij'
C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]
如果到目前为止您已经能够遵循该代码,那么恭喜您!这就是您需要阅读einsum方程式所需的全部。请特别注意原始einsum公式如何映射到以上代码段中的最终sumsum语句。for循环和范围边界只是模糊不清的,最终声明是您真正需要了解的所有内容。
为了完整起见,让我们看看如何确定每个范围变量的范围。嗯,每个变量的范围只是它索引的维度的长度。显然,如果变量在一个或多个张量中索引一个以上的维度,则每个维度的长度必须相等。这是上面带有完整范围的代码:
# C's shape is determined by the shapes of the inputs
# b indexes both A and B, so its range can come from either A.shape or B.shape
# i indexes only A, so its range can only come from A.shape, the same is true for j and B
assert A.shape[0] == B.shape[0]
assert A.shape[1] == B.shape[1]
assert A.shape[2] == B.shape[2]
C = np.zeros((A.shape[0], A.shape[3], B.shape[3]))
for b in range(A.shape[0]): # b indexes both A and B, or B.shape[0], which must be the same
for i in range(A.shape[3]):
for j in range(B.shape[3]):
# h and w can come from either A or B
for h in range(A.shape[1]):
for w in range(A.shape[2]):
C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]