2.6 PyTorch中张量的运算

2.6.1 涉及单个张量的函数运算

2.5节介绍了张量的一些基本属性操作。在进行深度学习的过程中，经常会遇到另一类操作——张量的运算，例如，对张量做四则运算、线性变换和激活等。这些基础操作既可以由张量自带的方法实现，也可以由torch包中的一些函数实现，具体如代码2.11所示。

代码2.11 PyTorch张量的函数运算。

从代码2.11中可以看到，对于大多数常用的函数，比如平方根函数sqrt，一般有两种调用方式，第一种是调用张量的内置方法（大多数的函数都有张量的内置方法版本，比如前面介绍过的sigmoid、relu和tanh等，以及常用的函数，如幂函数pow等），第二种是调用torch自带的函数。这两种操作的结果相同，均返回一个新张量，该张量的每个元素是原始张量的每个元素经过函数作用的结果。

前面也介绍过，很多张量的内置方法都有一个“下画线版本”，该版本的方法会直接改变调用方法的张量的值，这个操作也叫原地（In-Place）操作。一个常用的带下画线的方法是copy_方法，通过在这个方法中传入一个形状相同的张量，将会把这个张量的值复制到原张量中。需要注意的是，这里不能直接用等号进行赋值，因为Python的等号表示绑定变量和某个张量，如果直接用等号对张量赋值，得到的张量和等号右边的张量是同一个张量（copy_方法不会改变张量分配的内存地址，而是改变张量内存地址中存储的值）。对张量来说，也可以对自身的一些元素做四则运算，比如，经常用到的函数，包括求和的函数（和内置方法）torch.sum、求积的函数（和内置方法）torch.prod、求平均的函数（和内置方法）torch.mean。默认情况下，这些函数在进行求和、求积、求平均等计算的同时，会自动消除被计算的维度（即张量的维度被缩减），如果要保留这些维度，需要设置参数keepdim=True，这样这个维度就会被保留为1。

2.6.2 涉及多个张量的函数运算

除前面介绍的以一个张量作为参数的操作外，还有以两个张量作为参数的操作。比如，两个形状相同的张量之间逐个元素的四则运算（参与运算的两个元素的位置一一对应）。这里既可以使用加、减、乘、除的运算符进行张量间的运算，也可以使用add、sub、mul和div方法来进行运算。同样，这些内置方法有原地操作版本add_、sub_、mul_/div_，具体可以参考代码2.12。

代码2.12 PyTorch张量的四则运算。

2.6.3 张量的极值和排序

在进行深度学习的过程中，我们经常需要获取张量（沿着某个维度）的最大值和最小值，以及这些值所在的位置。如果只需要最大值或者最小值的位置，可以使用argmax和argmin，通过传入具体要沿着哪个维度求最大值和最小值的位置，返回沿着该维度最大和最小值对应的序号是多少。如果既要求获得最大和最小值的位置，又要求获得具体的值，就需要使用max和min，通过传入具体的维度，同时返回沿着该维度最大和最小值的位置，以及对应最大值和最小值组成的元组（Tuple）。

最后一个和大小有关的函数是排序函数sort（默认顺序是从小到大，如果要从大到小排序，需要设置参数descending=True），同样是传入具体需要进行排序的维度，返回的是排序完的张量，以及对应排序后的元素在原始张量上的位置。如果想知道原始张量的元素沿着某个维度排第几位，只需要对相应排序后的元素在原始张量上的位置进行再次排序，得到的新位置的值即为原始张量沿着该方向进行大小排序后的序号。同前面一样，关于极值和排序的函数，既可以是PyTorch的函数，也可以是张量的内置方法，两种方法调用的效果等价。具体函数的调用可参考代码2.13。

代码2.13 PyTorch极值和排序的函数。

2.6.4 矩阵的乘法和张量的缩并

除四则运算、最大和最小值，以及排序运算外，由两个张量作为参数的操作还有矩阵乘法（线性变换）。在Python 3.5以后，＠运算符号可以作为矩阵乘法的运算符号（参考Python的PEP 465标准）。因此，有几种方法来实现矩阵乘法。第一种是使用torch.mm函数来进行矩阵乘法（mm代表Matrix Multiplication），第二种是使用张量内置的mm方法来进行矩阵乘法，第三种是利用＠运算符号来实现，具体如代码2.14所示。

代码2.14 PyTorch张量的矩阵乘法运算。

另一个特殊的矩阵乘法的函数是bmm函数。在进行深度学习的过程中，实际经常用到的是迷你批次的数据，一般来说，第一个维度是（迷你）批次的大小。因此，数据的矩阵实际上是一个（迷你）批次的矩阵，即一个三维的张量（可以看作是一个迷你批次数量的矩阵叠加在一起）。在这种情况下，如果两个张量做矩阵乘法，一般情况下是沿着（迷你）批次的方向分别对每个矩阵对做乘法，最后把所有乘积的结果整合在一起。如果是大小为b×m×k的张量和b×k×n的张量相乘，那么结果应该是一个b×m×n的张量。

对于更大维度的张量的乘积，往往要决定各自张量元素乘积的结果需要沿着哪些维度求和，这个操作称为缩并（Contraction）。这时就需要引入爱因斯坦求和约定（Einstein Summation Convention），具体如公式2.1所示。

其中，参与运算的有两个张量，分别记为A和B，输出结果为C，这里把对应维度的下标分为三类：在A、B、C中都出现的，意味着这两个下标对应的一系列元素需要做两两乘积（即张量积）；如果在A、B中出现，但C中并没有出现，意味着这两个下标对应的一系列元素要做乘积求和（类似于内积）；在A、B中出现，C中只出现一次且这两个指标对应的维度大小相等，意味着这两个维度之间元素按照位置做乘法。

在上述条件下，前面介绍的矩阵乘法和批次矩阵乘法都能归结为爱因斯坦求和乘法，具体在PyTorch中对应的函数为torch.einsum。这里可以参考代码2.15来看一下这个函数的用法。

代码2.15 torch.einsum函数的使用。

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

2.6.5 张量的拼接和分割

在实际应用中，经常会碰到的另一种情况是把不同的张量按照某一个维度组合在一起，或者把一个张量按照一定的形状进行分割，这时候就需要用到张量的组合和分割函数，主要有以下几个函数：torch.stack、torch.cat、torch.split和torch.chunk。其中前两个函数负责将多个张量堆叠和拼接成一个张量，后两个函数负责把一个张量分割成多个张量。

● torch.stack函数的功能是通过传入的张量列表，同时指定并创建一个维度，把列表的张量沿着该维度堆叠起来，并返回堆叠以后的张量。传入的张量列表中所有张量的大小必须一致。

● torch.cat函数通过传入的张量列表指定某一个维度，把列表中的张量沿着该维度堆叠起来，并返回堆叠以后的张量。传入的张量列表的所有张量除指定堆叠的维度外，其他的维度大小必须一致。这个函数和torch.stack函数类似，都是对张量进行组合，两个函数的区别在于，前者的维度一开始并不存在，会新建一个维度，后者的维度则是预先存在的，所有的张量会沿着这个维度堆叠。

● torch.split函数的功能是执行前面叠加函数的反向操作，最后输出的是张量沿着某个维度分割后的列表。该函数需要传入三个参数，即被分割的张量、分割后维度的大小（整数或者列表）和分割的维度。如果传入整数，则沿着传入的维度分割成好几段，每段沿着该维度的大小是传入的整数（如果传入张量的维度不能被分割后的张量整除，则最后一个张量在该维度的大小会小于传入的整数）；如果传入整数列表，则按照列表整数的大小来分割这个维度。

● torch.chunk函数与torch.split函数的功能类似，区别在于前者传入的整数参数是分割的段数，输入张量在该维度的大小需要被分割的段数整除。另外，张量有内置的split和chunk方法，与torch.split和torch.chunk函数等价。

有关张量拼接和分割的代码可以参考代码2.16。

代码2.16 张量的拼接和分割。

2.6.6 张量维度的扩增和压缩

对于张量，还有一个比较常用的操作是沿着某个方向对张量做扩增（Expand）或对张量进行压缩（Squeeze）。这两种情况与张量的大小等于1的维度有关。对一个张量来说，可以任意添加一个维度，该维度的大小为1，而不改变张量的数据，因为张量的大小等于所有维度大小的乘积，那些为1的维度不改变张量的大小。于是我们就可以自由地在张量中添加任意数目为1的维度。在PyTorch中，使用torch.unsqueeze函数和张量unsqueeze方法来增加张量的维度，其中传入的参数为需要扩增的维度。同样，假如有一个张量的一些维度大小为1，就能直接压缩这些维度，具体使用的是torch.squeeze函数和张量的squeeze方法。

张量维度扩增和压缩的代码如代码2.17所示。

代码2.17 张量维度扩增和压缩。

2.6.7 张量的广播

张量的扩增有助于实现张量的另外一种功能，即张量的广播（Broadcast）。在张量的运算中会碰到一种情况，即两个不同维度张量之间做四则运算，且两个张量的某些维度相等。显然，如果按照张量的四则运算的定义，两个不同维度的张量不能进行四则运算。为了能够让它们进行计算，首先需要把维度数目比较小的张量扩增到和维度数目比较大的张量一致，这就需要使用unsqueeze方法来对张量进行维度扩增。完成扩增维度的两个张量必须能够在维度上对齐，即两个张量之间对应的维度存在两种情况，至少有一个维度大小为1，或者两个维度大小均不为1，但是相等。

下面举例来说，假设一个张量的大小为3×4×5，另外一个张量大小为3×5，为了能够让两个张量进行四则运算，需要把第二个张量的形状展开成3×1×5，这样两个张量就能对齐。关于大小为3×4×5的张量如何和大小为3×1×5的张量进行四则运算，其定义是将3×1×5的张量沿着第二个维度复制4次，使之成为3×4×5的张量，这样这两个张量就能进行元素一一对应的计算。

关于张量广播的具体代码可以参考代码2.18。

代码2.18 张量的广播。