机器学习 - Python - numpy.einsum()

机器学习记录

np.einsum()是numpy中的一个万能函数,可以进行广泛的张量积运算。它通过一系列缩写表示法,计算多维数组之间的相乘积。

import numpy as np

a = np.arange(25).reshape(5,5)
b = np.arange(5)
c = np.arange(6).reshape(2,3)

a

array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

b

array([0, 1, 2, 3, 4])

c

array([[0, 1, 2],
       [3, 4, 5]])

# 对角线之和
np.einsum('ii', a)

60

# 对角线之和
np.trace(a)

60

# 对角线之和
np.einsum(a, [0, 0])

60

# 对角线向量
np.einsum('ii->i', a)

array([ 0,  6, 12, 18, 24])

# 对角线向量
np.diag(a)

array([ 0,  6, 12, 18, 24])

# 对角线上元素 组成的向量
np.einsum(a, [0, 0], [0])

array([ 0,  6, 12, 18, 24])

# Sum over an axis (requires explicit form):
# 对一个轴求和（需要明确的形式）：
# 每行元素和 组成的向量
np.einsum('ij->i', a)

array([ 10,  35,  60,  85, 110])

np.einsum(a, [0,1], [0])

array([ 10,  35,  60,  85, 110])

np.sum(a, axis=1)

array([ 10,  35,  60,  85, 110])

np.einsum('...j->...', a)

array([ 10,  35,  60,  85, 110])

np.einsum(a, [Ellipsis,1], [Ellipsis])

array([ 10,  35,  60,  85, 110])

# 每列元素和 组成的向量
np.einsum('ij->j', a)

array([50, 55, 60, 65, 70])

np.einsum(a, [1,0], [0])

array([50, 55, 60, 65, 70])

np.sum(a, axis=0)

array([50, 55, 60, 65, 70])

# Compute a matrix transpose, or reorder any number of axes:
# 矩阵转置
np.einsum('ji', c)

array([[0, 3],
       [1, 4],
       [2, 5]])

np.einsum('ij->ji', c)

array([[0, 3],
       [1, 4],
       [2, 5]])

np.einsum(c, [1,0])

array([[0, 3],
       [1, 4],
       [2, 5]])

np.transpose(c)

array([[0, 3],
       [1, 4],
       [2, 5]])

# Vector inner products:
# 向量内积：
# 1. 向量, 2. 元素自己相乘和求和
np.einsum('i,i', b, b)

30

np.einsum(b, [0], b, [0])

30

np.inner(b,b)

30

# Matrix vector multiplication:
# 矩阵向量乘法：
np.einsum('ij,j', a, b)

array([ 30,  80, 130, 180, 230])

np.einsum(a, [0,1], b, [1])

array([ 30,  80, 130, 180, 230])

np.dot(a, b)

array([ 30,  80, 130, 180, 230])

np.einsum('...j,j', a, b)

array([ 30,  80, 130, 180, 230])

# Broadcasting and scalar multiplication:
# 矩阵乘以系数
np.einsum('..., ...', 3, a)

array([[ 0,  3,  6,  9, 12],
       [15, 18, 21, 24, 27],
       [30, 33, 36, 39, 42],
       [45, 48, 51, 54, 57],
       [60, 63, 66, 69, 72]])

np.einsum(',ij', 3, c)

array([[ 0,  3,  6],
       [ 9, 12, 15]])

np.einsum(3, [Ellipsis], c, [Ellipsis])

array([[ 0,  3,  6],
       [ 9, 12, 15]])

np.multiply(3, c)

array([[ 0,  3,  6],
       [ 9, 12, 15]])

# Vector outer product:
# 向量外积：
# a0*b0 + a0*b1 + a0*b2 +...
# a1*b0 + a1*b1 + a1*b2 +...
# ...
np.einsum('i,j', np.arange(2)+1, b)

array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])

np.einsum(np.arange(2)+1, [0], b, [1])

array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])

np.outer(np.arange(2)+1, b)

array([[0, 1, 2, 3, 4],
       [0, 2, 4, 6, 8]])

# Tensor contraction:
# 张量收缩：
a = np.arange(60.).reshape(3,4,5)
b = np.arange(24.).reshape(4,3,2)
np.einsum('ijk,jil->kl', a, b)

array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

np.einsum(a, [0,1,2], b, [1,0,3], [2,3])

array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

np.tensordot(a,b, axes=([1,0],[0,1]))

array([[4400., 4730.],
       [4532., 4874.],
       [4664., 5018.],
       [4796., 5162.],
       [4928., 5306.]])

# Writeable returned arrays
# 可写返回数组
# 设置 元素的值
a = np.zeros((3, 3))
np.einsum('ii->i', a)[:] = 1
a

array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]])

# Example of ellipsis use
# 省略号使用示例
a = np.arange(6).reshape((3,2))
b = np.arange(12).reshape((4,3))
np.einsum('ki,jk->ij', a, b)

array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

np.einsum('ki,...k->i...', a, b)

array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

np.einsum('k...,jk', a, b)

array([[10, 28, 46, 64],
       [13, 40, 67, 94]])

a = np.ones(64).reshape(2,4,8)
a

array([[[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]],

       [[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]]])

for iteration in range(500):
    _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a)
a

array([[[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]],

       [[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]]])

for iteration in range(500):
    _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')
a

array([[[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]],

       [[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]]])

for iteration in range(500):
    _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='greedy')
a

array([[[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]],

       [[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]]])

path = np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize='optimal')[0]
for iteration in range(500):
    _ = np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a, optimize=path)
a

array([[[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]],

       [[1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1., 1., 1., 1.]]])

官方文档