np。einsum工作吗?

给定数组A和B,它们的矩阵乘法和转置是使用(A @ B).T计算的,或者等效地,使用:

np.einsum("ij, jk -> ki", A, B)

当前回答

我认为最简单的例子是在tensorflow文档中

将方程转换为einsum符号有四个步骤。我们以这个方程为例C[i,k] = sum_j A[i,j] * B[j,k]

首先,我们删除变量名。得到ik = sum_j ij * jk 我们删除sum_j项,因为它是隐式的。得到ik = ij * jk 我们用,替换*。我们得到ik = ij jk 输出在RHS上,用->符号分隔。得到ij jk - >ik

einsum解释器只是反向运行这4个步骤。对结果中缺失的所有指标求和。

下面是文档中的更多示例

# Matrix multiplication
einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Trace
einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]

# Batch matrix multiplication
einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]

其他回答

关于np.einsum的另一个观点

这里的大多数答案都是通过例子来解释的,我想我会给出一个额外的观点。

你可以把einsum看成一个广义矩阵求和算子。给定的字符串包含下标,下标是表示轴的标签。我喜欢把它看作你的操作定义。下标提供了两个明显的约束条件:

每个输入数组的轴数, 输入之间的轴大小相等。

让我们以最初的例子为例:np。einsum('ij,jk->ki', A, B)。翻译成A.ndim == 2和B.ndim == 2,和2。to A.shape[1] == B.shape[0]。

稍后您将看到,还有其他限制条件。例如:

输出下标中的标签不能出现多次。 输出下标中的标签必须出现在输入下标中。

当看ij,jk->ki时,你可以把它想象成:

输入数组中的哪些组件将贡献给输出数组中的组件[k, i]。

下标包含输出数组的每个组件的操作的准确定义。

我们将继续使用操作ij,jk->ki,以及A和B的以下定义:

>>> A = np.array([[1,4,1,7], [8,1,2,2], [7,4,3,4]])
>>> A.shape
(3, 4)

>>> B = np.array([[2,5], [0,1], [5,7], [9,2]])
>>> B.shape
(4, 2)

输出Z的形状为(B.shape[1], a .shape[0]),可以简单地按以下方式构造。从Z的空数组开始:

Z = np.zeros((B.shape[1], A.shape[0]))
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        for k range(B.shape[0]):
           Z[k, i] += A[i, j]*B[j, k] # ki <- ij*jk

np。Einsum是关于在输出数组中积累贡献。每个(A[i,j], B[j,k])对都对每个Z[k, i]分量有贡献。

你可能已经注意到,它看起来非常类似于你计算一般矩阵乘法的方式……


最小的实现

这里是np的一个最小实现。Python中的einsum。这应该有助于理解在引擎盖下面到底发生了什么。

在我们继续进行的过程中,我将继续引用前面的例子。将输入定义为[A, B]。

np。Einsum实际上可以接受两个以上的输入。在下面,我们将关注一般情况:n个输入和n个输入下标。主要目标是找到迭代的定义域,即所有值域的笛卡尔积。

我们不能依赖于手动编写for循环,因为我们不知道会有多少个循环。主要思想是这样的:我们需要找到所有唯一的标签(我将使用key和keys来引用它们),找到相应的数组形状,然后为每个数组创建范围,并使用itertools计算范围的乘积。产品得到了研究领域。

index keys constraints sizes ranges
1 'i' A.shape[0] 3 range(0, 3)
2 'j' A.shape[1] == B.shape[0] 4 range(0, 4)
0 'k' B.shape[1] 2 range(0, 2)

研究的领域是笛卡尔积:范围(0,2)x范围(0,3)x范围(0,4)。

Subscripts processing: >>> expr = 'ij,jk->ki' >>> qry_expr, res_expr = expr.split('->') >>> inputs_expr = qry_expr.split(',') >>> inputs_expr, res_expr (['ij', 'jk'], 'ki') Find the unique keys (labels) in the input subscripts: >>> keys = set([(key, size) for keys, input in zip(inputs_expr, inputs) for key, size in list(zip(keys, input.shape))]) {('i', 3), ('j', 4), ('k', 2)} We should be checking for constraints (as well as in the output subscript)! Using set is a bad idea but it will work for the purpose of this example. Get the associated sizes (used to initialize the output array) and construct the ranges (used to create our domain of iteration): >>> sizes = dict(keys) {'i': 3, 'j': 4, 'k': 2} >>> ranges = [range(size) for _, size in keys] [range(0, 2), range(0, 3), range(0, 4)] We need an list containing the keys (labels): >>> to_key = sizes.keys() ['k', 'i', 'j'] Compute the cartesian product of the ranges >>> domain = product(*ranges) Note: [itertools.product][1] returns an iterator which gets consumed over time. Initialize the output tensor as: >>> res = np.zeros([sizes[key] for key in res_expr]) We will be looping over domain: >>> for indices in domain: ... pass For each iteration, indices will contain the values on each axis. In our example, that would provide i, j, and k as a tuple: (k, i, j). For each input (A and B) we need to determine which component to fetch. That's A[i, j] and B[j, k], yes! However, we don't have variables i, j, and k, literally speaking. We can zip indices with to_key to create a mapping between each key (label) and its current value: >>> vals = dict(zip(to_key, indices)) To get the coordinates for the output array, we use vals and loop over the keys: [vals[key] for key in res_expr]. However, to use these to index the output array, we need to wrap it with tuple and zip to separate the indices along each axis: >>> res_ind = tuple(zip([vals[key] for key in res_expr])) Same for the input indices (although there can be several): >>> inputs_ind = [tuple(zip([vals[key] for key in expr])) for expr in inputs_expr] We will use a itertools.reduce to compute the product of all contributing components: >>> def reduce_mult(L): ... return reduce(lambda x, y: x*y, L) Overall the loop over the domain looks like: >>> for indices in domain: ... vals = {k: v for v, k in zip(indices, to_key)} ... res_ind = tuple(zip([vals[key] for key in res_expr])) ... inputs_ind = [tuple(zip([vals[key] for key in expr])) ... for expr in inputs_expr] ... ... res[res_ind] += reduce_mult([M[i] for M, i in zip(inputs, inputs_ind)])


>>> res
array([[70., 44., 65.],
       [30., 59., 68.]]) 

这很接近np。einsum('ij,jk->ki', A, B)返回!

我发现NumPy: The tricks of The trade (Part II)很有教育意义

我们使用->表示输出数组的顺序。把'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求和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]])
>>>

我认为最简单的例子是在tensorflow文档中

将方程转换为einsum符号有四个步骤。我们以这个方程为例C[i,k] = sum_j A[i,j] * B[j,k]

首先,我们删除变量名。得到ik = sum_j ij * jk 我们删除sum_j项,因为它是隐式的。得到ik = ij * jk 我们用,替换*。我们得到ik = ij jk 输出在RHS上,用->符号分隔。得到ij jk - >ik

einsum解释器只是反向运行这4个步骤。对结果中缺失的所有指标求和。

下面是文档中的更多示例

# Matrix multiplication
einsum('ij,jk->ik', m0, m1)  # output[i,k] = sum_j m0[i,j] * m1[j, k]

# Dot product
einsum('i,i->', u, v)  # output = sum_i u[i]*v[i]

# Outer product
einsum('i,j->ij', u, v)  # output[i,j] = u[i]*v[j]

# Transpose
einsum('ij->ji', m)  # output[j,i] = m[i,j]

# Trace
einsum('ii', m)  # output[j,i] = trace(m) = sum_i m[i, i]

# Batch matrix multiplication
einsum('aij,ajk->aik', s, t)  # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]

在阅读einsum方程时,我发现能够这样做最有帮助 在心里把它们归结为它们的命令版本。

让我们从下面的陈述开始:

C = np.einsum('bhwi,bhwj->bij', A, B)

首先看标点符号,我们看到箭头前面有两个4个逗号分隔的斑点——bhwi和bhwj, 后面还有一个3个字母的bij。因此,该方程由两个秩4张量输入产生一个秩3张量结果。

现在,让每个blob中的每个字母都是一个范围变量的名称。字母出现在圆点中的位置 是它在这个张量中所覆盖的轴的下标。 因此,生成C的每个元素的命令式求和必须从三个嵌套的for循环开始,每个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] = ...

所以,本质上,对于c的每个输出索引都有一个for循环,我们暂时不确定范围。

接下来我们看左边,有没有范围变量没有出现在右边?在我们的例子中,是的,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公式如何映射到上面代码片段中的最终求和语句。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]

让我们创建两个数组,它们具有不同但兼容的尺寸,以突出它们的相互作用

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的最后一个,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,:,:]

也就是说,在a的末尾增加一个k维,在B的前面增加一个i维,得到一个(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]