文章最前： 我是Octopus，这个名字来源于我的中文名–章鱼；我热爱编程、热爱算法、热爱开源。所有源码在我的个人github ；这博客是记录我学习的点点滴滴，如果您对 Python、Java、AI、算法有兴趣，可以关注我的动态，一起学习，共同进步。

4-2,张量的数学运算

张量数学运算主要有：标量运算，向量运算，矩阵运算，以及 使用非常强大而灵活的爱因斯坦求和函数torch.einsum进行任意维的张量运算。

此外我们还会介绍张量运算的广播机制。

本篇文章内容如下：

• 标量运算

• 向量运算

• 矩阵运算

• 任意维张量运算(torch.einsum)

• 广播机制

本节中的torch.einsum的理解是重难点。

import torch 
print("torch.__version__="+torch.__version__) 

torch.__version__=2.0.1

一，标量运算 (操作的张量至少是0维)

张量的数学运算符可以分为标量运算符、向量运算符、以及矩阵运算符。

加减乘除乘方，以及三角函数，指数，对数等常见函数，逻辑比较运算符等都是标量运算符。

标量运算符的特点是对张量实施逐元素运算。

有些标量运算符对常用的数学运算符进行了重载。并且支持类似numpy的广播特性。

import torch 
import numpy as np 

a = torch.tensor(1.0)
b = torch.tensor(2.0)
a+b 

tensor(3.)

a = torch.tensor([[1.0,2],[-3,4.0]])
b = torch.tensor([[5.0,6],[7.0,8.0]])
a+b  #运算符重载

tensor([[ 6.,  8.],
        [ 4., 12.]])

a-b 

tensor([[ -4.,  -4.],
        [-10.,  -4.]])

a*b 

tensor([[  5.,  12.],
        [-21.,  32.]])

a/b

tensor([[ 0.2000,  0.3333],
        [-0.4286,  0.5000]])

a**2

tensor([[ 1.,  4.],
        [ 9., 16.]])

a**(0.5)

tensor([[1.0000, 1.4142],
        [   nan, 2.0000]])

a%3 #求模

tensor([[1., 2.],
        [-0., 1.]])

torch.div(a, b, rounding_mode='floor')  #地板除法

tensor([[ 0.,  0.],
        [-1.,  0.]])

a>=2 # torch.ge(a,2)  #ge: greater_equal缩写

tensor([[False,  True],
        [False,  True]])

(a>=2)&(a<=3)

tensor([[False,  True],
        [False, False]])

(a>=2)|(a<=3)

tensor([[True, True],
        [True, True]])

a==5 #torch.eq(a,5)

tensor([[False, False],
        [False, False]])

torch.sqrt(a)

tensor([[1.0000, 1.4142],
        [   nan, 2.0000]])

a = torch.tensor([1.0,8.0])
b = torch.tensor([5.0,6.0])
c = torch.tensor([6.0,7.0])

d = a+b+c
print(d)

tensor([12., 21.])

print(torch.max(a,b))

tensor([5., 8.])

print(torch.min(a,b))

tensor([1., 6.])

x = torch.tensor([2.6,-2.7])

print(torch.round(x)) #保留整数部分，四舍五入
print(torch.floor(x)) #保留整数部分，向下归整
print(torch.ceil(x))  #保留整数部分，向上归整
print(torch.trunc(x)) #保留整数部分，向0归整

tensor([ 3., -3.])
tensor([ 2., -3.])
tensor([ 3., -2.])
tensor([ 2., -2.])

x = torch.tensor([2.6,-2.7])
print(torch.fmod(x,2)) #作除法取余数 
print(torch.remainder(x,2)) #作除法取剩余的部分，结果恒正

tensor([ 0.6000, -0.7000])
tensor([0.6000, 1.3000])

# 幅值裁剪
x = torch.tensor([0.9,-0.8,100.0,-20.0,0.7])
y = torch.clamp(x,min=-1,max = 1)
z = torch.clamp(x,max = 1)
print(y)
print(z)

tensor([ 0.9000, -0.8000,  1.0000, -1.0000,  0.7000])
tensor([  0.9000,  -0.8000,   1.0000, -20.0000,   0.7000])

relu = lambda x:x.clamp(min=0.0)
relu(torch.tensor(5.0))

tensor(5.)

二，向量运算 （原则上操作的张量至少是一维张量）

向量运算符只在一个特定轴上运算，将一个向量映射到一个标量或者另外一个向量。

#统计值

a = torch.arange(1,10).float().view(3,3)

print(torch.sum(a))
print(torch.mean(a))
print(torch.max(a))
print(torch.min(a))
print(torch.prod(a)) #累乘
print(torch.std(a))  #标准差
print(torch.var(a))  #方差
print(torch.median(a)) #中位数

tensor(45.)
tensor(5.)
tensor(9.)
tensor(1.)
tensor(362880.)
tensor(2.7386)
tensor(7.5000)
tensor(5.)

#指定维度计算统计值
b = torch.arange(1,10).float().view(3,3)
print(b)
print(torch.max(b,dim = 0))
print(torch.max(b,dim = 1))

tensor([[1., 2., 3.],
        [4., 5., 6.],
        [7., 8., 9.]])
torch.return_types.max(
values=tensor([7., 8., 9.]),
indices=tensor([2, 2, 2]))
torch.return_types.max(
values=tensor([3., 6., 9.]),
indices=tensor([2, 2, 2]))

#cum扫描
a = torch.arange(1,10)

print(torch.cumsum(a,0))
print(torch.cumprod(a,0))
print(torch.cummax(a,0).values)
print(torch.cummax(a,0).indices)
print(torch.cummin(a,0))

tensor([ 1,  3,  6, 10, 15, 21, 28, 36, 45])
tensor([     1,      2,      6,     24,    120,    720,   5040,  40320, 362880])
tensor([1, 2, 3, 4, 5, 6, 7, 8, 9])
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])
torch.return_types.cummin(
values=tensor([1, 1, 1, 1, 1, 1, 1, 1, 1]),
indices=tensor([0, 0, 0, 0, 0, 0, 0, 0, 0]))

#torch.sort和torch.topk可以对张量排序
a = torch.tensor([[9,7,8],[1,3,2],[5,6,4]]).float()
print(torch.topk(a,2,dim = 0),"\n")
print(torch.topk(a,2,dim = 1),"\n")
print(torch.sort(a,dim = 1),"\n")

#利用torch.topk可以在Pytorch中实现KNN算法

torch.return_types.topk(
values=tensor([[9., 7., 8.],
        [5., 6., 4.]]),
indices=tensor([[0, 0, 0],
        [2, 2, 2]])) 

torch.return_types.topk(
values=tensor([[9., 8.],
        [3., 2.],
        [6., 5.]]),
indices=tensor([[0, 2],
        [1, 2],
        [1, 0]])) 

torch.return_types.sort(
values=tensor([[7., 8., 9.],
        [1., 2., 3.],
        [4., 5., 6.]]),
indices=tensor([[1, 2, 0],
        [0, 2, 1],
        [2, 0, 1]])) 

三，矩阵运算 （操作的张量至少是二维张量）

矩阵必须是二维的。类似torch.tensor([1,2,3])这样的不是矩阵。

矩阵运算包括：矩阵乘法，矩阵逆，矩阵求迹，矩阵范数，矩阵行列式，矩阵求特征值，矩阵分解等运算。

#矩阵乘法
a = torch.tensor([[1,2],[3,4]])
b = torch.tensor([[2,0],[0,2]])
print(a@b)  #等价于torch.matmul(a,b) 或 torch.mm(a,b)

tensor([[2, 4],
        [6, 8]])

#高维张量的矩阵乘法在后面的维度上进行
a = torch.randn(5,5,6)
b = torch.randn(5,6,4)
(a@b).shape 

torch.Size([5, 5, 4])

#矩阵转置
a = torch.tensor([[1.0,2],[3,4]])
print(a.t())

tensor([[1., 3.],
        [2., 4.]])

#矩阵逆，必须为浮点类型
a = torch.tensor([[1.0,2],[3,4]])
print(torch.inverse(a))

tensor([[-2.0000,  1.0000],
        [ 1.5000, -0.5000]])

#矩阵求trace
a = torch.tensor([[1.0,2],[3,4]])
print(torch.trace(a))

tensor(5.)

#矩阵求范数
a = torch.tensor([[1.0,2],[3,4]])
print(torch.norm(a))

tensor(5.4772)

#矩阵行列式
a = torch.tensor([[1.0,2],[3,4]])
print(torch.det(a))

tensor(-2.)

#矩阵特征值和特征向量
a = torch.tensor([[1.0,2],[-5,4]],dtype = torch.float)
print(torch.linalg.eig(a))

#两个特征值分别是 -2.5+2.7839j, 2.5-2.7839j 

torch.return_types.linalg_eig(
eigenvalues=tensor([2.5000+2.7839j, 2.5000-2.7839j]),
eigenvectors=tensor([[0.2535-0.4706j, 0.2535+0.4706j],
        [0.8452+0.0000j, 0.8452-0.0000j]]))

#矩阵QR分解, 将一个方阵分解为一个正交矩阵q和上三角矩阵r
#QR分解实际上是对矩阵a实施Schmidt正交化得到q

a  = torch.tensor([[1.0,2.0],[3.0,4.0]])
q,r = torch.linalg.qr(a)
print(q,"\n")
print(r,"\n")
print(q@r)

tensor([[-0.3162, -0.9487],
        [-0.9487,  0.3162]]) 

tensor([[-3.1623, -4.4272],
        [ 0.0000, -0.6325]]) 

tensor([[1.0000, 2.0000],
        [3.0000, 4.0000]])

#矩阵svd分解
#svd分解可以将任意一个矩阵分解为一个正交矩阵u,一个对角阵s和一个正交矩阵v.t()的乘积
#svd常用于矩阵压缩和降维
a=torch.tensor([[1.0,2.0],[3.0,4.0],[5.0,6.0]])

u,s,v = torch.linalg.svd(a)

print(u,"\n")
print(s,"\n")
print(v,"\n")

import torch.nn.functional as F 
print(u@F.pad(torch.diag(s),(0,0,0,1))@v.t())

#利用svd分解可以在Pytorch中实现主成分分析降维

tensor([[-0.2298,  0.8835,  0.4082],
        [-0.5247,  0.2408, -0.8165],
        [-0.8196, -0.4019,  0.4082]]) 

tensor([9.5255, 0.5143]) 

tensor([[-0.6196, -0.7849],
        [-0.7849,  0.6196]]) 

tensor([[1.0000, 2.0000],
        [3.0000, 4.0000],
        [5.0000, 6.0000]])

四，任意维张量运算(torch.einsum)

如果问pytorch中最强大的一个数学函数是什么？

我会说是torch.einsum：爱因斯坦求和函数。

它几乎是一个"万能函数"：能实现超过一万种功能的函数。

不仅如此，和其它pytorch中的函数一样，torch.einsum是支持求导和反向传播的，并且计算效率非常高。

einsum 提供了一套既简洁又优雅的规则，可实现包括但不限于：内积，外积，矩阵乘法，转置和张量收缩（tensor contraction）等张量操作，熟练掌握 einsum 可以很方便的实现复杂的张量操作，而且不容易出错。

尤其是在一些包括batch维度的高阶张量的相关计算中，若使用普通的矩阵乘法、求和、转置等算子来实现很容易出现维度匹配等问题，若换成einsum则会特别简单。

套用一句深度学习paper标题当中非常时髦的话术，einsum is all you needed 😋！

1，einsum规则原理

顾名思义，einsum这个函数的思想起源于家喻户晓的小爱同学：爱因斯坦~。

很久很久以前，小爱同学在捣鼓广义相对论。广义相对论表述各种物理量用的都是张量。

比如描述时空有一个四维时空度规张量，描述电磁场有一个电磁张量，描述运动的有能量动量张量。

在理论物理学家中，小爱同学的数学基础不算特别好，在捣鼓这些张量的时候，他遇到了一个比较头疼的问题：公式太长太复杂了。

有没有什么办法让这些张量运算公式稍微显得对人类友好一些呢，能不能减少一些那种扭曲的$\sum$求和符号呢？

小爱发现，求和导致维度收缩，因此求和符号操作的指标总是只出现在公式的一边。

例如在我们熟悉的矩阵乘法中

$$C_{ij}=\sum _k{A}_{ik}{B}_{kj}$$

k这个下标被求和了，求和导致了这个维度的消失，所以它只出现在右边而不出现在左边。

这种只出现在张量公式的一边的下标被称之为哑指标，反之为自由指标。

小爱同学脑瓜子一转，反正这种只出现在一边的哑指标一定是被求和求掉的，干脆把对应的$\sum$求和符号省略得了。

这就是爱因斯坦求和约定：

只出现在公式一边的指标叫做哑指标，针对哑指标的$\sum$求和符号可以省略。

公式立刻清爽了很多。

$$C_{ij}=A_{ik}{B}_{kj}$$

这个公式表达的含义如下：

C这个张量的第i行第j列由$A$这个张量的第i行第k列和$B$这个张量的第k行第j列相乘，这样得到的是一个三维张量$D$, 其元素为$D_{ikj}$，然后对$D$在维度k上求和得到。

公式展现形式中除了省去了求和符号，还省去了乘法符号(代数通识)。

借鉴爱因斯坦求和约定表达张量运算的清爽整洁，numpy、tensorflow和 torch等库中都引入了 einsum这个函数。

上述矩阵乘法可以被einsum这个函数表述成

C = torch.einsum("ik,kj->ij",A,B)

这个函数的规则原理非常简洁，3句话说明白。

• 1，用元素计算公式来表达张量运算。

• 2，只出现在元素计算公式箭头左边的指标叫做哑指标。

• 3，省略元素计算公式中对哑指标的求和符号。

import torch 

A = torch.tensor([[1,2],[3,4.0]])
B = torch.tensor([[5,6],[7,8.0]])

C1 = A@B
print(C1)

C2 = torch.einsum("ik,kj->ij",[A,B])
print(C2)

tensor([[19., 22.],
        [43., 50.]])
tensor([[19., 22.],
        [43., 50.]])

2，einsum基础范例

einsum这个函数的精髓实际上是第一条:

用元素计算公式来表达张量运算。

而绝大部分张量运算都可以用元素计算公式很方便地来表达，这也是它为什么会那么神通广大。

#例1，张量转置
A = torch.randn(3,4,5)

#B = torch.permute(A,[0,2,1])
B = torch.einsum("ijk->ikj",A) 

print("before:",A.shape)
print("after:",B.shape)

before: torch.Size([3, 4, 5])
after: torch.Size([3, 5, 4])

#例2，取对角元
A = torch.randn(5,5)
#B = torch.diagonal(A)
B = torch.einsum("ii->i",A)
print("before:",A.shape)
print("after:",B.shape)

before: torch.Size([5, 5])
after: torch.Size([5])

#例3，求和降维
A = torch.randn(4,5)
#B = torch.sum(A,1)
B = torch.einsum("ij->i",A)
print("before:",A.shape)
print("after:",B.shape)

before: torch.Size([4, 5])
after: torch.Size([4])

#例4，哈达玛积
A = torch.randn(5,5)
B = torch.randn(5,5)
#C=A*B
C = torch.einsum("ij,ij->ij",A,B)
print("before:",A.shape, B.shape)
print("after:",C.shape)

before: torch.Size([5, 5]) torch.Size([5, 5])
after: torch.Size([5, 5])

#例5，向量内积
A = torch.randn(10)
B = torch.randn(10)
#C=torch.dot(A,B)
C = torch.einsum("i,i->",A,B)
print("before:",A.shape, B.shape)
print("after:",C.shape)

before: torch.Size([10]) torch.Size([10])
after: torch.Size([])

#例6，向量外积(类似笛卡尔积)
A = torch.randn(10)
B = torch.randn(5)
#C = torch.outer(A,B)
C = torch.einsum("i,j->ij",A,B)
print("before:",A.shape, B.shape)
print("after:",C.shape)

before: torch.Size([10]) torch.Size([5])
after: torch.Size([10, 5])

#例7，矩阵乘法
A = torch.randn(5,4)
B = torch.randn(4,6)
#C = torch.matmul(A,B)
C = torch.einsum("ik,kj->ij",A,B)
print("before:",A.shape, B.shape)
print("after:",C.shape)

before: torch.Size([5, 4]) torch.Size([4, 6])
after: torch.Size([5, 6])

#例8，张量缩并
A = torch.randn(3,4,5)
B = torch.randn(4,3,6)
#C = torch.tensordot(A,B,dims=[(0,1),(1,0)])
C = torch.einsum("ijk,jih->kh",A,B)
print("before:",A.shape, B.shape)
print("after:",C.shape)

before: torch.Size([3, 4, 5]) torch.Size([4, 3, 6])
after: torch.Size([5, 6])

3，einsum高级范例

einsum可用于超过两个张量的计算。

例如：双线性变换。这是向量内积的一种扩展，一种常用的注意力机制实现方式)

不考虑batch维度时，双线性变换的公式如下:

$$A=qW{k}^T$$

考虑batch维度时，无法用矩阵乘法表示，可以用元素计算公式表达如下：

$$A_{ij}=\sum _k\sum _l{Q}_{ik}{W}_{jkl}{K}_{il}=Q_{ik}{W}_{jkl}{K}_{il}$$

#例9，bilinear注意力机制

#====不考虑batch维度====
q = torch.randn(10) #query_features
k = torch.randn(10) #key_features
W = torch.randn(5,10,10) #out_features,query_features,key_features
b = torch.randn(5) #out_features

#a = q@[email protected]()+b  
a = torch.bilinear(q,k,W,b)
print("a.shape:",a.shape)


#=====考虑batch维度====
Q = torch.randn(8,10)    #batch_size,query_features
K = torch.randn(8,10)    #batch_size,key_features
W = torch.randn(5,10,10) #out_features,query_features,key_features
b = torch.randn(5)       #out_features

#A = torch.bilinear(Q,K,W,b)
A = torch.einsum('bq,oqk,bk->bo',Q,W,K) + b
print("A.shape:",A.shape)

a.shape: torch.Size([5])
A.shape: torch.Size([8, 5])

我们也可以用einsum来实现更常见的scaled-dot-product 形式的 Attention.

不考虑batch维度时，scaled-dot-product形式的Attention用矩阵乘法公式表示如下:

$$a=softmax(\frac{q{k}^T}{d_k})$$

考虑batch维度时，无法用矩阵乘法表示，可以用元素计算公式表达如下：

$$A_{ij}=softmax(\frac{Q_{in}{K}_{ijn}}{d_k})$$

#例10，scaled-dot-product注意力机制

#====不考虑batch维度====
q = torch.randn(10)  #query_features
k = torch.randn(6,10) #key_size, key_features

d_k = k.shape[-1]
a = torch.softmax(q@k.t()/d_k,-1) 

print("a.shape=",a.shape )

#====考虑batch维度====
Q = torch.randn(8,10)  #batch_size,query_features
K = torch.randn(8,6,10) #batch_size,key_size,key_features

d_k = K.shape[-1]
A = torch.softmax(torch.einsum("in,ijn->ij",Q,K)/d_k,-1) 

print("A.shape=",A.shape )

a.shape= torch.Size([6])
A.shape= torch.Size([8, 6])

#性能测试

#=====考虑batch维度====
Q = torch.randn(80,100)    #batch_size,query_features
K = torch.randn(80,100)    #batch_size,key_features
W = torch.randn(50,100,100) #out_features,query_features,key_features
b = torch.randn(50)       #out_features

%%timeit 
A = torch.bilinear(Q,K,W,b)

1.83 ms ± 78.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit 
A = torch.einsum('bq,oqk,bk->bo',Q,W,K) + b

636 µs ± 27.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

五，广播机制

Pytorch的广播规则和numpy是一样的:

• 1、如果张量的维度不同，将维度较小的张量进行扩展，直到两个张量的维度都一样。
• 2、如果两个张量在某个维度上的长度是相同的，或者其中一个张量在该维度上的长度为1，那么我们就说这两个张量在该维度上是相容的。
• 3、如果两个张量在所有维度上都是相容的，它们就能使用广播。
• 4、广播之后，每个维度的长度将取两个张量在该维度长度的较大值。
• 5、在任何一个维度上，如果一个张量的长度为1，另一个张量长度大于1，那么在该维度上，就好像是对第一个张量进行了复制。

torch.broadcast_tensors可以将多个张量根据广播规则转换成相同的维度。

维度扩展允许的操作有两种：
1，增加一个维度
2，对长度为1的维度进行复制扩展

a = torch.tensor([1,2,3])
b = torch.tensor([[0,0,0],[1,1,1],[2,2,2]])
print(b + a) 

tensor([[1, 2, 3],
        [2, 3, 4],
        [3, 4, 5]])

torch.cat([a[None,:]]*3,dim=0) + b  

tensor([[1, 2, 3],
        [2, 3, 4],
        [3, 4, 5]])

a_broad,b_broad = torch.broadcast_tensors(a,b)
print(a_broad,"\n")
print(b_broad,"\n")
print(a_broad + b_broad) 

tensor([[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]]) 

tensor([[0, 0, 0],
        [1, 1, 1],
        [2, 2, 2]]) 

tensor([[1, 2, 3],
        [2, 3, 4],
        [3, 4, 5]])