4-2张量的数学运算

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

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

本篇文章内容如下：

• 标量运算
• 向量运算
• 矩阵运算
• 任意维张量运算
• 广播机制

import torch

print("torch.__version__="+torch.__version__)
"""
torch.__version__=2.1.1+cu118
"""

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)
"""
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.)
"""

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

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

# 统计值
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))
print(torch.topk(a, 2, dim=1))
print(torch.sort(a, dim=1))
"""
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]]))
"""

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

矩阵必须是二维的。类似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(a.det())  # torch.det(a)
"""
tensor(-2.)
"""

# 矩阵特征值和特征向量
a = torch.tensor([[1.0, 2], [-5, 4]], dtype=torch.float)
print(torch.linalg.eig(a))
"""
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分解实际上是对矩阵q实施Schmidt正交化得到q
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
q, r = torch.linalg.qr(a)
print(q)
print(r)
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)
print(s)
print(v)

import torch.nn.functional as F

print(u@F.pad(torch.diag(s), (0, 0, 0, 1))@v.t())  # 左右上下，下侧填充0
"""
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]])
"""

4.任意维张量运算（torch.einsum）

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

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

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

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

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

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

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

• 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.]])
"""

• einsum基础范例

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

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

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

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

# B = torch.permute(A, [0, 2, 1])
B = torch.einsum('ijk->ikj', A)
print(A.shape)
print(B.shape)
"""
torch.Size([3, 4, 5])
torch.Size([3, 5, 4])
"""

# 例2，取对角元
A = torch.randn(5, 5)
# B = torch.diagonal(A)
B = torch.einsum('ii->i', A)
print(A.shape)
print(B.shape)
"""
torch.Size([5, 5])
torch.Size([5])
"""

# 例3，求和降维
A = torch.randn(4, 5)
# B = torch.sum(A, 1)
B = torch.einsum('ij->i', A)
print(A.shape)
print(B.shape)
"""
torch.Size([4, 5])
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(A.shape, B.shape)
print(C.shape)
"""
torch.Size([5, 5]) torch.Size([5, 5])
torch.Size([5, 5])
"""

# 例5，向量内积
A = torch.randn(10)
B = torch.randn(10)
# C = torch.dot(A, B)
C = torch.einsum('i,j->', A, B)
print(A.shape, B.shape)
print(C.shape)
"""
torch.Size([10]) torch.Size([10])
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(A.shape, B.shape)
print(C.shape)
"""
torch.Size([10]) torch.Size([5])
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(A.shape, B.shape)
print(C.shape)
"""
torch.Size([5, 4]) torch.Size([4, 6])
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(A.shape, B.shape)
print(C.shape)
"""
torch.Size([3, 4, 5]) torch.Size([4, 3, 6])
torch.Size([5, 6])
"""

• eimsum高级范例

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

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

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

\[A = qWk^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@W@k.t()+b  
a = torch.bilinear(q,k,W,b)
print("a.shape:",a.shape)
"""
a.shape: torch.Size([5])
"""

# 考虑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)
"""
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)
"""
torch.Size([6])
"""

# 考虑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)
"""
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)
"""
5.92 ms ± 2.13 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
"""

%%timeit
A = torch.einsum('bk,oqk,bk->bo', Q, W, K) + b
"""
1.2 ms ± 16.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
"""

5.广播机制

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

• 如果张量的维度不同，将维度较小的张量进行扩展，直到两个张量的维度都一致
• 如果两个张量在某个维度上的长度是相同的，或者其中一个张量在该维度上的长度是1，那么我们就说这两个张量在该维度上是相同的
• 如果量张量在所有维度上都是相容的，它们就能使用广播
• 广播之后，每个维度的长度将取两个张量在该维度长度的较大值
• 在任何一个维度上，如果一个张量的长度为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)
print(b_broad)
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]])
"""