pytorch中张量的运算

在深度学习中，经常需要对张量做四则、线性变换和激活等。

1. 单个张量的函数运算

t = torch.rand(2, 2)  # 产生一个3x3的张量
print(t)
print(t.sqrt())  # 张量的平方根，张量内部方法
print(torch.sqrt(t))  # 张量的平方根，函数形式
print(t.sqrt_())  # 平方根原地操作
print(torch.sum(t))  # 所有元素求和
print(torch.sum(t, 0))  # 第0维元素求和
print(torch.sum(t, [0, 1]))  # 对0,1维元素求和
print(torch.mean(t, [0, 1]))  # 对0,1维元素求平均

2. 多个张量的函数运算

pytorch中可以使用加减乘除的运算符进行张量间的运算，也可以使用add、sub、mul和div方法来进行运算。

t1 = torch.rand(2, 2)
t2 = torch.rand(2, 2)
# 元素相加
print(t1.add(t2))
print(t1+t2)
print('=' * 50)
# 元素相减
print(t1.sub(t2))
print(t1-t2)
print('=' * 50)
# 元素相乘
print(t1.mul(t2))
print(t1*t2)
print('=' * 50)
# 元素相除
print(t1.div(t2))
print(t1/t2)
print('=' * 50)

3. 张量的极值和排序

t = torch.randn(3, 3)
print(t)
print('=' * 50)
print(torch.max(t, 0))  # 沿着第0个维度的最大值
print(torch.argmax(t, 0))  # 沿着第0个维度的最大值索引位置
print('=' * 50)
print(torch.min(t, 0))  # 沿着第0个维度的最小值
print(torch.argmin(t, 0))  # 沿着第0个维度的最小值索引位置
print('=' * 50)
print(t.sort(-1))  # 沿着最后一个维度排序，返回排序后的张量和张量元素在改维度的原始位置

4. 矩阵的乘法和张量的缩并

4.1 pytorch张量的矩阵乘法运算

torch进行矩阵乘法可以用torch.mm（mm代表Matrix Multiplication），也可以用内置的mm方法来进行矩阵乘法，或者使用@运算符。

t1 = torch.randn(3, 4)  # 建立一个3x4的张量
t2 = torch.randn(4, 3)  # 建立一个4x3的张量
# 矩阵乘法
print(torch.mm(t1, t2))  # 矩阵乘法
print(t1.mm(t2))
print(t1@t2)
print('=' * 50)
a1 = torch.randn(2, 3, 4)
a2 = torch.randn(2, 4, 3)
print(torch.bmm(a1, a2))  # （迷你）批次矩阵乘法，返回结果为2x3x3
print(a1.bmm(a2))
print(a1@a2)

4.2 torch.einsum函数的使用

在计算较大维度的乘积时，通常需要决定各自张量元素乘法的结果沿着哪些维度求和，这个操作称为缩并，需要引入爱因斯坦求和约定（Einstein Summation Covention）。

$$c_{ijkabc}=\sum _{l_1{l}_2...l_n}=A_{ijk...m{l}_1{l}_2...l_n}{B}_{abc...m{l}_1{l}_2...l_n}$$
其中，参与运算的有两个张量，分别记为A和B，输出结果为C，这里把对应的维度下标分为三类：在A、B、C中都出现，意味着这两个下标对应的一系列元素做两两乘积（即张量积）；如果在A、B中出现，但C中并没有出现，意味着两个下标对应的一些列元素做乘积求和（类似于内积）；在A、B中出现，C中只出现一次且这两个指标对应的维度大小相等，意味着着两个维度之间按照位置做乘法。

t1 = torch.randn(2, 3, 4)
t2 = torch.randn(2, 4, 3)
print(t1.bmm(t2))  # 批次矩阵乘法的结果
print('=' * 50)
print(torch.einsum('bnk,bkl -> bnl', t1, t2))  # einsum计算

torch.einsum函数在使用的时候需要传入两个张量的下标对应的形状，以不同的字母区分（字母可以任意选择，只需要服从前面的规则即可），以及最后的形状，用->符号连接，最后传入两个输入的张量，即可得到输出的结果。需要注意的是，求和的指标所在维度的大小一定要相同，否则报错。

5. 张量的拼接和分割

t1 = torch.randn(3, 4)
t2 = torch.randn(3, 4)
t3 = torch.randn(3, 4)
print(torch.stack([t1,t2,t3], -1).shape)  #沿着最后一个维度做堆叠，返回大小为3x4x3的张量
print('=' * 50)
t = torch.randn(3, 6)
print(t.split([1,2,3], -1))  # 把张量沿着最后一个维度分割为三个张量
print('=' * 50)
print(t.split(3, -1))  # 把张良沿着最后一个维度分割，分割大小为3，输出的张量大小均为3x2
print('=' * 50)
print(t.chunk(3, -1))  # 把张良沿着最后一个维度分割，分割大小为3，输出的张量大小均为3x2

6. 张量维度的扩增和压缩

t = torch.rand(3, 4)
print(t.shape)
print(t.unsqueeze(-1).shape)  # 扩增最后一个维度
print(t.unsqueeze(-1).unsqueeze(-1).shape)  # 继续扩增最后一个维度
print('=' * 50)
t2 = torch.rand(1,3,4,1)
print(t.shape)
print(t.squeeze().shape)  # 压缩所有大小为1的维度

7. 张量的广播

假设一个张量的大小为3x4x5，另外一个张量大小为3x5，为了能够让着两个张量进行四则运算，就要把第二个张量的形状展开成3x1x5，这样两个张量就能对其。在两个张量进行计算的时候，3x1x5的张量会沿着第二个维度复制4次，使之成为3x4x5的张量，这样两个张量就能进行运算了。

t1 = torch.randn(3,4,5)
t2 = torch.randn(3,5)
t2 = t2.unsqueeze(1)  # 把张量的形状变为3x1x5
t3 = t1 + t2  # 广播求和，最后的结果为3x4x5的张量
print(t3)
print(t3.shape)