python 笔记:爱因斯坦求和 einsum

Posted UQI-LIUWJ

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了python 笔记:爱因斯坦求和 einsum相关的知识,希望对你有一定的参考价值。

1 einsum简介

        使用爱因斯坦求和约定,可以以简单的方式表示许多常见的多维线性代数数组运算。

        给定两个矩阵A和B,我们想对它们做一些操作,比如 multiply、sum或者transpose等。虽然numpy里面有可以直接使用的接口,能够实现这些功能,但是使用enisum可以做的更快、更节省空间。

        举例说明,我们现在有两个矩阵A和B。我们想计算A和B的哈达玛乘积(即逐元素乘积),然后按行求和。

import numpy as np
A = np.array([0, 1, 2])
B = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])

        如果我们不适用einsum的话,也不是不能计算,就是需要写几步来完成:

A.reshape(-1,1)
'''
array([[0],
       [1],
       [2]])
'''

A.reshape(-1,1)*B
'''
array([[ 0,  0,  0,  0],
       [ 4,  5,  6,  7],
       [16, 18, 20, 22]])
'''


(A.reshape(-1,1)*B).sum(axis=1)
#array([ 0, 22, 76])

        如果我们使用einsum的话,一行就可以实现:

np.einsum('i,ij->i', A, B)
#array([ 0, 22, 76])

  2 einsum原理

                使用einsum的关键是,正确地labelling(标记)输入数组和输出数组的axes(轴)。

                我们可以使用字符串(比如:ijk,这种表示方式更常用)或者一个整数列表(比如:[0,1])来标记axes。

                比如,为了实现矩阵乘法,我们可以用einsum这么写(至于为什么这个是矩阵乘法,我们在后面会说明)

np.einsum('ij,jk->ik', A, B)

                字符串'ij,jk->ik'可以根据'->'的位置来切分,左边的部分('ij,jk')标记了输入的axes,右边的('ik')标记了输出的axes。

                输入标记又根据','的位置进行切分,'ij'标记了第一个输入A的axes,'jk'标记了第二个输入B的axes。

                'ij'、'jk'的字符长度都是2,对应着A和B为2D数组,'ik'的长度也为2,因此输出也是2D数组。

                给定输入

A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])

         np.einsum('ij,jk->ik', A, B)可以看作是:

  • 在输入数组的标记之间,重复字母表示沿这些轴的值将相乘,这些乘积构成输出数组的值。比如图中沿着j轴做乘积。
  • 从输出标记中省略的字母表示沿该轴的值将被求和。比如图中的输出没有包含j轴,因此沿着j轴求和得到了输出数组中的每一项。
    • 如果输出的标记是'ijk',那么会得到一个 3x3x3 的矩阵。 

    A = np.array([[1, 3, 5],
                  [7, 9, -7],
                  [-5, -3, -1]])
    B = np.array([[0, 2,4],
                  [6, 8, 6],
                  [4, 2, 0]])
    np.einsum('ij,jk->ijk', A, B)
    
    
    '''
    array([[[  0,   2,   4],
            [ 18,  24,  18],
            [ 20,  10,   0]],
    
           [[  0,  14,  28],
            [ 54,  72,  54],
            [-28, -14,   0]],
    
           [[  0, -10, -20],
            [-18, -24, -18],
            [ -4,  -2,   0]]])
    '''

 输出标记是'ik'的时候,并不会创建中间的 3x3x3 的矩阵,而是直接将总和累加到2D数组中。

A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])
np.einsum('ij,jk->ik', A, B)
'''
array([[ 38,  36,  22],
       [ 26,  72,  82],
       [-22, -36, -38]])
'''

如果输出的标记是空,那么输出整个矩阵的和

A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])
np.einsum('ij,jk->', A, B)
#180

 我们可以按任意顺序排序不求和的轴。

A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])
np.einsum('ij,jk->kji', A, B)

'''
array([[[  0,   0,   0],
        [ 18,  54, -18],
        [ 20, -28,  -4]],

       [[  2,  14, -10],
        [ 24,  72, -24],
        [ 10, -14,  -2]],

       [[  4,  28, -20],
        [ 18,  54, -18],
        [  0,   0,   0]]])
'''

3 einsum分析

3.1 'ij,jk->ijk'  与 'ij,jk->kji'

 我们一个一个分析一下

A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])
np.einsum('ij,jk->ijk', A, B)


'''
array([[[  0,   2,   4],
        [ 18,  24,  18],
        [ 20,  10,   0]],

       [[  0,  14,  28],
        [ 54,  72,  54],
        [-28, -14,   0]],

       [[  0, -10, -20],
        [-18, -24, -18],
        [ -4,  -2,   0]]])
'''

首先,这几个数字是怎么得到的?

0=1*02=1*24=1*4
18=3*624=3*818=3*6
20=5*410=5*20=5*0
0=7*014=7*228=7*4
54=9*672=9*854=9*6
-28=-7*4-14=-7*20=-7*0
0=-5*0-10=-5*2-20=-5*4
-18=-3*6-24=-3*8-18=-3*6
-4=-1*4-2=-1*20=-1*0

转换成坐标,有:

[0,0]*[0,0][0,0]*[0,1][0,0]*[0,2]
[0,1]*[1,0][0,1]*[1,1][0,1]*[1,2]
[0,2]*[2,0][0,2]*[2,1][0,2]*[2,2]
[1,0]*[0,0]

[1,0]*[0,2]

[1,0]*[0,4]
[1,1]*[1,0][1,1]*[1,1][1,1]*[1,2]
[1,2]*[2,0][1,2]*[2,1][1,2]*[2,2]
[2,0]*[0,0][2,0]*[0,1][2,0]*[0,2]
[2,1]*[1,0][2,1]*[1,1][2,1]*[1,2]
[2,2]*[2,0][2,2]*[2,1][2,2]*[2,2]
A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])
np.einsum('ij,jk->kji', A, B)

'''
array([[[  0,   0,   0],
        [ 18,  54, -18],
        [ 20, -28,  -4]],

       [[  2,  14, -10],
        [ 24,  72, -24],
        [ 10, -14,  -2]],

       [[  4,  28, -20],
        [ 18,  54, -18],
        [  0,   0,   0]]])
'''

 与上面类似,我们就看第一个3*3的矩阵吧

0=1*00=7*00=-5*0
18=3*654=9*6-18=-3*6
20=5*4-28=-7*4-4=-1*4
[0,0]*[0,0][1,0]*[0,0][2,0]*[0,0]
[0,1]*[1,0][1,1]*[1,0][2,1]*[1,0]
[0.2]*[2,0][1,2]*[2,0][2,2]*[2,0]

可以这么考虑 对于 结果矩阵(比如ijk),第【i,j,k】元素的结果等于【i,j】乘以【j,k】

3.2 'ij,jk->ik'  

A = np.array([[1, 3, 5],
              [7, 9, -7],
              [-5, -3, -1]])
B = np.array([[0, 2,4],
              [6, 8, 6],
              [4, 2, 0]])
np.einsum('ij,jk->ik', A, B)
'''
array([[ 38,  36,  22],
       [ 26,  72,  82],
       [-22, -36, -38]])
'''

我们前面 'ij,jk->ijk'的结果是 

0=1*02=1*24=1*4
18=3*624=3*818=3*6
20=5*410=5*20=5*0
0=7*014=7*228=7*4
54=9*672=9*854=9*6
-28=-7*4-14=-7*20=-7*0
0=-5*0-10=-5*2-20=-5*4
-18=-3*6-24=-3*8-18=-3*6
-4=-1*4-2=-1*20=-1*0

这边相当于

        

38=0+18+2036=2+24+1022=4+18
26=54-2872=14+72-1482=54+28
-22=-18-4-36=-10-24-2-38=-20-18

可以这么考虑 对于 结果矩阵(比如ik),第【i,k】元素的结果等于:对所有的j,【i,j】乘以【j,k】的结果的和

4 常用的Einsum

4.1 向量篇

('i',A)

向量A的一个视图

('i->', A)

sum(A)

('i,i->i', A,B)

向量A,B对应位置相乘

'i,i->i':结果的第i位,是A和B的第i位的积 

('i,i->', A,B)

向量A,B的内积

 'i,i->':结果的第i位,是A和B的第i位的积 载求和 

('i,j->ij', A,B)

向量A,B的外积

 'i,j->ij':结果的第i行第j列,是A的第i个元素和B的第j个元素的乘积

 4.2 矩阵篇

('ij', A)

返回矩阵A

('ji->ij', A)

返回矩阵A的转置

结果的第i行第j列的元素是A的第j行第i列的元素 

('ii->i', A)

矩阵A的对角线元素

结果的第i个元素是A的第i行第i列的元素 

('ij->', A)

矩阵A的元素之和

结果是A的第i行第j列的元素,再求和 

('ij->j', A)

A纵向求和

结果的第j个元素是,对所有的i,A的第(i,j)个元素的和 

('ij->i', A)

A横向求和

结果的第i个元素是,对所有的j,A的第(i,j)个元素的和 

('ij,ij->ij', A,B)

矩阵A,B相应位置的乘积

结果第i,j个元素,是A的第(i,j)个元素和B的第(i,j)个元素的乘积 

('ij,ji->ij', A,B)

矩阵A和  矩阵B的转置   相应位置的乘积

结果第i,j个元素,是A的第(i,j)个元素和B的第(j,i)个元素的乘积 

('ij,jk->ik', A,B)

A,B的矩阵乘积

结果的第(i,k)个元素等于A的第(i,j)个元素乘以B的第(j,k)个元素 

('ij,kj->ik', A,B)

矩阵A和矩阵B的内积

('ij,kl->jikl', A,B)

A的每个元素乘以矩阵B

结果的 第(j,i,k,l)个元素是A的(i,j)和B的(k,l)的乘积 

 5 显示表明和隐式表明

我们将指定'->'和输出标记称为 explicit mode。

如果不指定'->'和输出标记,numpy会将输入标记中只出现一次的标记按照字母表顺序,作为输出标记(也就是 implicit mode)。

'ij,jk->ik' 等价于 'ij,jk'

参考文章:einsum初探 - 知乎 (zhihu.com)

以上是关于python 笔记:爱因斯坦求和 einsum的主要内容,如果未能解决你的问题,请参考以下文章

Python 模块 Cupy 函数在使用 cupy.einsum() 时出错

哈希算法爱因斯坦求和约定,这是2020年的注意力机制

斯坦福CS231n—深度学习与计算机视觉----学习笔记 课时6

`numpy.einsum` 中的 `out` 参数无法按预期工作

PyTorch 中的 tensordot 以及 einsum 函数介绍

RuntimeError: einsum(): operands do not broadcast with remapped shapes [original->remapped]