Pytorch学习笔记（一）—— Pytorch基础

前导

更多文章代码详情可查看博主个人网站：https://www.iwtmbtly.com/

导入需要使用的库和文件：

>>> import torch
>>> import numpy as np

一、什么是Tensor

在深度学习中，从数据的组织，到模型内部的参数，都是通过一种叫做张量（Tensor）的数据结构进行表示和处理。

Tensor 是深度学习框架中极为基础的概念，也是 PyTroch、TensorFlow 中最重要的知识点之一，它是一种数据的存储和处理结构。

回忆一下我们目前知道的几种数据表示：

标量，也称 Scalar，是一个只有大小，没有方向的量，比如 1.8、e、10 等。
向量，也称 Vector，是一个有大小也有方向的量，比如 (1,2,3,4) 等。
矩阵，也称 Matrix，是多个向量合并在一起得到的量，比如[(1,2,3),(4,5,6)]等。

为了帮助你更好理解标量、向量和矩阵，我特意准备了一张示意图，可以结合图片理解。

不难发现，几种数据表示其实都是有着联系的，标量可以组合成向量，向量可以组合成矩阵。那么，我们可否将它们看作是一种数据形式呢？

答案是可以的，这种统一的数据形式，在 PyTorch 中我们称之为张量 (Tensor)。从标量、向量和矩阵的关系来看，你可能会觉得它们就是不同“维度”的 Tensor，这个说法对，也不全对。

说它不全对是因为在 Tensor 的概念中，我们更愿意使用 Rank（秩）来表示这种“维度”，比如标量，就是 Rank 为 0 阶的 Tensor；向量就是 Rank 为 1 阶的 Tensor；矩阵就是 Rank 为 2 阶的 Tensor。也有 Rank 大于 2 的 Tensor。当然啦，你如果说维度其实也没什么错误，平时很多人也都这么叫。

说完 Tensor 的含义，我们一起看一下 Tensor 的类型，以及如何创建 Tensor。

二、Tensor 的类型、创建及转换

在不同的深度学习框架下，Tensor 呈现的特点大同小异，我们使用它的方法也差不多。这节课我们就以 PyTorch 中的使用方法为例进行学习。

Tensor 的类型

在 PyTorch 中，Tensor 支持的数据类型有很多种，这里列举较为常用的几种格式：

一般来说，torch.float32、torch.float64、torch.uint8 和 torch.int64 用得相对较多一些，但是也不是绝对，还是要根据实际情况进行选择。

Tensor 的创建

PyTorch 对于 Tensor 的操作已经非常友好了，你可以通过多种不同的方式创建一个任意形状的 Tensor，而且每种方式都很简便，我们一起来看一下。

直接创建

首先来看直接创建的方法，这也是最简单创建的方法。我们需要用到下面的 torch.tensor 函数直接创建。

torch.tensor(data, dtype=None, device=None,requires_grad=False)

例如：

>>> A = torch.tensor([x for x in range(6)])
>>> A
tensor([0, 1, 2, 3, 4, 5])

结合代码，我们看看其中的参数是什么含义。

我们从左往右依次来看，首先是 data，也就是我们要传入模型的数据。PyTorch 支持通过 list、 tuple、numpy array、scalar 等多种类型进行数据传入，并转换为 tensor。

接着是 dtype，它声明了你需要返回一个怎样类型的 Tensor，具体类型可以参考前面表格里列举的 Tensor 的那些类型。

然后是 device，这个参数指定了数据要返回到的设备，目前暂时不需要关注，缺省即可。

最后一个参数是 requires_grad，用于说明当前量是否需要在计算中保留对应的梯度信息。在 PyTorch 中，只有当一个 Tensor 设置 requires_grad 为 True 的情况下，才会对这个 Tensor 以及由这个 Tensor 计算出来的其他 Tensor 进行求导，然后将导数值存在 Tensor 的 grad 属性中，便于优化器来更新参数。

所以，你需要注意的是，把 requires_grad 设置成 true 或者 false 要灵活处理。如果是训练过程就要设置为 true，目的是方便求导、更新参数。而到了验证或者测试过程，我们的目的是检查当前模型的泛化能力，那就要把 requires_grad 设置成 Fasle，避免这个参数根据 loss 自动更新。

从 NumPy 中创建

在实际应用中，我们在处理数据的阶段很多时候使用的是 NumPy，而数据处理好之后想要传入 PyTorch 的深度学习模型中，则需要借助 Tensor，所以 PyTorch 提供了一个从 NumPy 转到 Tensor 的语句：

torch.from_numpy(ndarry)

例如：

>>> A = torch.from_numpy(np.array([x for x in range(6)]))
>>> A
tensor([0, 1, 2, 3, 4, 5])

有时候我们在开发模型的过程中，需要用到一些特定形式的矩阵 Tensor，比如全是 0 的，或者全是 1 的。这时我们就可以用这个方法创建，比如说，先生成一个全是 0 的 NumPy数组，然后转换成 Tensor。但是这样也挺麻烦的，因为这意味着你要引入更多的包（NumPy），也会使用更多的代码，这会增加出错的可能性。

不过你别担心，PyTorch 内部已经提供了更为简便的方法，我们接着往下看。

创建特殊形式的 Tensor

我们一块来看一下后面的几个常用函数，它们都是在 PyTorch 模型内部使用的。

创建零矩阵 Tensor：零矩阵顾名思义，就是所有的元素都为 0 的矩阵。

torch.zeros(size, dtype=None...)

例如：

>>> A = torch.zeros((3,4), dtype=float)
>>> A
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]], dtype=torch.float64)

其中，我们用得比较多的就是 size 参数和 dtype 参数。size 定义输出张量形状的整数序列。

这里你可能注意到了，在函数参数列表中我加入了省略号，这意味着 torch.zeros 的参数有很多。不过，现在是介绍零矩阵的概念，形状相对来说更重要。其他的参数（比如前面提到的 requires_grad 参数）与此无关，现阶段我们暂时不关注。

创建单位矩阵 Tensor：单位矩阵是指主对角线上的元素都为 1 的矩阵。

torch.eye(size, dtype=None...)

例如：

>>> A = torch.eye(6)
>>> A
tensor([[1., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0.],
        [0., 0., 0., 0., 1., 0.],
        [0., 0., 0., 0., 0., 1.]])

创建全一矩阵 Tensor：全一矩阵顾名思义，就是所有的元素都为 1 的矩阵。

torch.ones(size, dtype=None...)

例如：

>>> A = torch.ones((6,6))
>>> A
tensor([[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.]])

创建随机矩阵 Tensor：在 PyTorch 中有几种较为经常使用的随机矩阵创建方式，分别如下。

torch.rand(size)
torch.randn(size)
torch.normal(size, mean, std)
torch.randint(low, high, size）

这些方式各自有不同的用法，你可以根据自己的需要灵活使用。

torch.rand 用于生成数据类型为浮点型且维度指定的随机 Tensor，随机生成的浮点数据在 0~1 区间均匀分布。
torch.randn 用于生成数据类型为浮点型且维度指定的随机 Tensor，随机生成的浮点数的取值满足均值为 0、方差为 1 的标准正态分布。
torch.normal 用于生成数据类型为浮点型且维度指定的随机 Tensor，可以指定均值和标准差。
torch.randint 用于生成随机整数的 Tensor，其内部填充的是在[low,high) 均匀生成的随机整数。

Tensor 的转换

在实际项目中，我们接触到的数据类型有很多，比如 Int、list、NumPy 等。为了让数据在各个阶段畅通无阻，不同数据类型与 Tensor 之间的转换就非常重要了。接下来我们一起来看看 int、list、NumPy 是如何与 Tensor 互相转换的。

int 与 Tensor 的转换：

>>> a = torch.Tensor([1])
>>> b = a.item()
>>> b
1.0

我们通过 torch.Tensor 将一个数字（或者标量）转换为 Tensor，又通过 item() 函数，将 Tensor 转换为数字（标量），item() 函数的作用就是将 Tensor 转换为一个 python number。

list 与 tensor 的转换：

>>> a = [1, 2, 3]
>>> b = torch.Tensor(a)
>>> c = b.numpy().tolist()
>>> c
[1.0, 2.0, 3.0]

在这里对于一个 list a，我们仍旧直接使用 torch.Tensor，就可以将其转换为 Tensor 了。而还原回来的过程要多一步，需要我们先将 Tensor 转为 NumPy 结构，之后再使用 tolist() 函数得到 list。

NumPy 与 Tensor 的转换：

有了前面两个例子，你是否能想到 NumPy 怎么转换为 Tensor 么？对，我们仍旧 torch.Tensor 即可，是不是特别方便。

CPU 与 GPU 的 Tensor 之间的转换：

CPU->GPU: data.cuda()
GPU->CPU: data.cpu()

三、Tensor 的常用操作

好，刚才我们一起了解了 Tensor 的类型，如何创建 Tensor，以及如何实现 Tensor 和一些常见的数据类型之间的相互转换。其实 Tensor 还有一些比较常用的功能，比如获取形状、维度转换、形状变换以及增减维度，接下来我们一起来看看这些功能。

获取形状

在深度学习网络的设计中，我们需要时刻对 Tensor 的情况做到了如指掌，其中就包括获取 Tensor 的形式、形状等。

为了得到 Tensor 的形状，我们可以使用 shape 或 size 来获取。两者的不同之处在于， shape 是 torch.tensor 的一个属性，而 size() 则是一个 torch.tensor 拥有的方法。

>>> a = torch.zeros(2, 3, 5)
>>> a.shape
torch.Size([2, 3, 5])
>>> a.size()
torch.Size([2, 3, 5])

知道了 Tensor 的形状，我们就能知道这个 Tensor 所包含的元素的数量了。具体的计算方法就是直接将所有维度的大小相乘，比如上面的 Tensor a 所含有的元素的个数为 235=30 个。这样似乎有点麻烦，我们在 PyTorch 中可以使用 numel() 函数直接统计元素数量。

>>> a.numel()
30

矩阵转置(维度转换）

在 PyTorch 中有两个函数，分别是 permute() 和 transpose() 可以用来实现矩阵的转秩，或者说交换不同维度的数据。比如在调整卷积层的尺寸、修改 channel 的顺序、变换全连接层的大小的时候，我们就要用到它们。

其中，用 permute 函数可以对任意高维矩阵进行转置，但只有 tensor.permute() 这个调用方式，我们先看一下代码：

>>> x = torch.rand(2,3,5)
>>> x.shape
torch.Size([2, 3, 5])
>>> x = x.permute(2,1,0)
>>> x.shape
torch.Size([5, 3, 2])

有没有发现，原来的 Tensor 的形状是[2,3,5]，我们在 permute 中分别写入原来索引位置的新位置，x.permute(2,1,0)，2 表示原来第二个维度现在放在了第零个维度；同理 1 表示原来第一个维度仍旧在第一个维度；0 表示原来第 0 个维度放在了现在的第 2 个维度，形状就变成了[5,3,2]。

而另外一个函数 transpose，不同于 permute，它每次只能转换两个维度，或者说交换两个维度的数据。我们还是来看一下代码：

>>> x.shape
torch.Size([5, 3, 2])
>>> x = x.transpose(1,0)
>>> x.shape
torch.Size([3, 5, 2])

需要注意的是，经过了 transpose 或者 permute 处理之后的数据，变得不再连续了，什么意思呢？

还是接着刚才的例子说，我们使用 torch.rand(2,3,4) 得到的 tensor，在内存中是连续的，但是经过 transpose 或者 permute 之后呢，比如 transpose(1,0)，内存虽然没有变化，但是我们得到的数据“看上去”是第 0 和第 1 维的数据发生了交换，现在的第 0 维是原来的第 1 维，所以 Tensor 都会变得不再连续。

那你可能会问了，不连续就不连续呗，好像也没啥影响吧？这么想你就草率了，我们继续来看看 Tensor 的形状变换，学完以后你就知道 Tensor 不连续的后果了。

形状变换

在 PyTorch 中有两种常用的改变形状的函数，分别是 view 和 reshape。我们先来看一下 view。

>>> x = torch.randn(4, 4)
>>> x.shape
torch.Size([4, 4])
>>> x = x.view(2,8)
>>> x.shape
torch.Size([2, 8])

我们先声明了一个[4, 4]大小的 Tensor，然后通过 view 函数，将其修改为[2, 8]形状的 Tensor。我们还是继续刚才的 x，再进行一步操作，代码如下：

>>> x = x.permute(1,0)
>>> x.shape
torch.Size([8, 2])
>>> x.view(4, 4)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

结合代码可以看到，利用 permute，我们将第 0 和第 1 维度的数据进行了变换，得到了 [8, 2]形状的 Tensor，在这个新 Tensor 上进行 view 操作，忽然就报错了，为什么呢？其实就是因为 view 不能处理内存不连续 Tensor 的结构。

那这时候要怎么办呢？我们可以使用另一个函数，reshape：

>>> x = x.reshape(4, 4)
>>> x.shape
torch.Size([4, 4])

增减维度

有时候我们需要对 Tensor 增加或者删除某些维度，比如删除或者增加图片的几个通道。 PyTorch 提供了 squeeze() 和 unsqueeze() 函数解决这个问题。

我们先来看 squeeze()。如果 dim 指定的维度的值为 1，则将该维度删除，若指定的维度值不为 1，则返回原来的 Tensor。为了方便你理解，我还是结合例子来讲解。

>>> x = torch.rand(2,1,3)
>>> x.shape
torch.Size([2, 1, 3])
>>> y = x.squeeze(1)
>>> y.shape
torch.Size([2, 3])
>>> z = y.squeeze(1)
>>> z.shape
torch.Size([2, 3])

结合代码我们可以看到，我们新建了一个维度为[2, 1, 3]的 Tensor，然后将第 1 维度的数据删除，得到 y，squeeze 执行成功是因为第 1 维度的大小为 1。然而在 y 上我们打算进一步删除第 1 维度的时候，就会发现删除失败了，这是因为 y 此刻的第 1 维度的大小为 3，suqeeze 不能删除。

unsqueeze()：这个函数主要是对数据维度进行扩充。给指定位置加上维数为 1 的维度，我们同样结合代码例子来看看。

>>> x = torch.rand(2,1,3)
>>> y = x.unsqueeze(2)
>>> y.shape
torch.Size([2, 1, 1, 3])

这里我们新建了一个维度为[2, 1, 3]的 Tensor，然后在第 2 维度插入一个数据，这样就得到了一个[2,1,1,3]大小的 tensor。

四、Tensor 的连接操作

在项目开发中，深度学习某一层神经元的数据可能有多个不同的来源，那么就需要将数据进行组合，这个组合的操作，我们称之为连接。

cat

连接的操作函数如下：

torch.cat(tensors, dim = 0, out = None)

cat 是 concatnate 的意思，也就是拼接、联系的意思。该函数有两个重要的参数需要你掌握。

第一个参数是 tensors，它很好理解，就是若干个我们准备进行拼接的 Tensor。

第二个参数是 dim，我们回忆一下 Tensor 的定义，Tensor 的维度（秩）是有多种情况的。比如有两个 3 维的 Tensor，可以有几种不同的拼接方式（如下图），dim 参数就可以对此作出约定。

看到这里，你可能觉得上面画的图是三维的，看起来比较晦涩，所以咱们先从简单的二维的情况说起，我们先声明两个 3x3 的矩阵，代码如下：

>>> A=torch.ones(3,3)
>>> B=2*torch.ones(3,3)
>>> A
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
>>> B
tensor([[2., 2., 2.],
        [2., 2., 2.],
        [2., 2., 2.]])

我们先看看 dim=0 的情况，拼接的结果是怎样的：

>>> C=torch.cat((A,B),0)
>>> C
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [2., 2., 2.],
        [2., 2., 2.],
        [2., 2., 2.]])

你会发现，两个矩阵是按照“行”的方向拼接的。

我们接下来再看看，dim=1 的情况是怎样的：

>>> D=torch.cat((A,B),1)
>>> D
tensor([[1., 1., 1., 2., 2., 2.],
        [1., 1., 1., 2., 2., 2.],
        [1., 1., 1., 2., 2., 2.]])

显然，两个矩阵，是按照“列”的方向拼接的。那如果 Tensor 是三维甚至更高维度的呢？其实道理也是一样的，dim 的数值是多少，两个矩阵就会按照相应维度的方向链接两个 Tensor。

看到这里你可能会问了，cat 实际上是将多个 Tensor 在已有的维度上进行连接，那如果想增加新的维度进行连接，又该怎么做呢？这时候就需要 stack 函数登场了。

stack

为了让你加深理解，我们还是结合具体例子来看看。假设我们有两个二维矩阵 Tensor，把它们“堆叠”放在一起，构成一个三维的 Tensor，如下图：

这相当于原来的维度（秩）是 2，现在变成了 3，变成了一个立体的结构，增加了一个维度。你需要注意的是，这跟前面的 cat 不同，cat 中示意图的例子，原来就是 3 维的，cat 之后仍旧是 3 维的，而现在咱们是从 2 维变成了 3 维。

在实际图像算法开发中，咱们有时候需要将多个单通道 Tensor（2 维）合并，得到多通道的结果（3 维）。而实现这种增加维度拼接的方法，我们把它叫做 stack。

stack 函数的定义如下：

torch.stack(inputs, dim=0)

其中，inputs 表示需要拼接的 Tensor，dim 表示新建立维度的方向。

那 stack 如何使用呢？我们一块来看一个例子：

>>> A=torch.arange(0,4)
>>> A
tensor([0, 1, 2, 3])
>>> B=torch.arange(5,9)
>>> B
tensor([5, 6, 7, 8])
>>> C=torch.stack((A,B),0)
>>> C
tensor([[0, 1, 2, 3],
        [5, 6, 7, 8]])
>>> D=torch.stack((A,B),1)
>>> D
tensor([[0, 5],
        [1, 6],
        [2, 7],
        [3, 8]])

结合代码，我们可以看到，首先我们构建了两个 4 元素向量 A 和 B，它们的维度是 1。然后，我们在 dim=0，也就是“行”的方向上新建一个维度，这样维度就成了 2，也就得到了 C。而对于 D，我们则是在 dim=1，也就是“列”的方向上新建维度。

五、Tensor 的切分操作

学完了连接操作之后，我们再来看看连接的逆操作：切分。

切分就是连接的逆过程，有了刚才的经验，你很容易就会想到，切分的操作也应该有很多种，比如切片、切块等。没错，切分的操作主要分为三种类型：chunk、split、unbind。

乍一看有不少，其实是因为它们各有特点，适用于不同的使用情景，让我们一起看一下。

chunk

chunk 的作用就是将 Tensor 按照声明的 dim，进行尽可能平均的划分。

比如说，我们有一个 32channel 的特征，需要将其按照 channel 均匀分成 4 组，每组 8 个 channel，这个切分就可以通过 chunk 函数来实现。具体函数如下：

torch.chunk(input, chunks, dim=0)

我们挨个来看看函数中涉及到的三个参数：

首先是 input，它表示要做 chunk 操作的 Tensor。

接着，我们看下 chunks，它代表将要被划分的块的数量，而不是每组的数量。请注意， chunks 必须是整型。

最后是 dim，想想这个参数是什么意思呢？对，就是按照哪个维度来进行 chunk。

还是跟之前一样，我们通过几个代码例子直观感受一下。我们从一个简单的一维向量开始：

>>> A=torch.tensor([x for x in range(1,11)])
>>> B = torch.chunk(A, 2, 0)
>>> B
(tensor([1, 2, 3, 4, 5]), tensor([ 6,  7,  8,  9, 10]))

这里我们通过 chunk 函数，将原来 10 位长度的 Tensor A，切分成了两个一样 5 位长度的向量。（注意，B 是两个切分结果组成的 tuple）。

那如果 chunk 参数不能够整除的话，结果会是怎样的呢？我们接着往下看：

>>> B = torch.chunk(A, 3, 0)
>>> B
(tensor([1, 2, 3, 4]), tensor([5, 6, 7, 8]), tensor([ 9, 10]))

我们发现，10 位长度的 Tensor A，切分成了三个向量，长度分别是 4，4，2 位。这是怎么分的呢，不应该是 3，3，4 这样更为平均的方式么？

想要解决问题，就得找到规律。让我们再来看一个更大一点的例子，将 A 改为 17 位长度。

>>> A=torch.tensor([x for x in range(1, 18)])
>>> B = torch.chunk(A, 4, 0)
>>> B
(tensor([1, 2, 3, 4, 5]), tensor([ 6,  7,  8,  9, 10]), tensor([11, 12, 13, 14, 15]), tensor([16, 17]))

17 位长度的 Tensor A，切分成了四个分别为 5，5，5，2 位长度的向量。这时候你就会发现，其实在计算每个结果元素个数的时候，chunk 函数是先做除法，然后再向上取整得到每组的数量。

比如上面这个例子，17/4=4.25，向上取整就是 5，那就先逐个生成若干个长度为 5 的向量，最后不够的就放在一块，作为最后一个向量（长度 2）。

那如果 chunk 参数大于 Tensor 可以切分的长度，又要怎么办呢？我们实际操作一下，代码如下：

>>> A=torch.tensor([1,2,3])
>>> B = torch.chunk(A, 5, 0)
>>> B
(tensor([1]), tensor([2]), tensor([3]))

显然，被切分的 Tensor 只能分成若干个长度为 1 的向量。

由此可以推论出二维的情况，我们再举一个例子，看看二维矩阵 Tensor 的情况：

>>> A=torch.ones(4,4)
>>> A
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> B = torch.chunk(A, 2, 0)
>>> B
(tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.]]), 
 tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.]]))

还是跟前面的 cat 一样，这里的 dim 参数，表示的是第 dim 维度方向上进行切分。

刚才介绍的 chunk 函数，是按照“切分成确定的份数”来进行切分的，那如果想按照“每份按照确定的大小”来进行切分，该怎样做呢？PyTorch 也提供了相应的方法，叫做 split。

split

split 的函数定义如下，跟前面一样，我们还是分别看看这里涉及的参数。

torch.split(tensor, split_size_or_sections, dim=0)

首先是 tensor，也就是待切分的 Tensor。

然后是 split_size_or_sections 这个参数。当它为整数时，表示将 tensor 按照每块大小为这个整数的数值来切割；当这个参数为列表时，则表示将此 tensor 切成和列表中元素一样大小的块。

最后同样是 dim，它定义了要按哪个维度切分。

同样的，我们举几个例子来看一下 split 的具体操作。首先是 split_size_or_sections 是整数的情况：

>>> A=torch.rand(4,4)
>>> A
tensor([[0.3480, 0.4864, 0.2071, 0.6849],
        [0.5214, 0.6468, 0.9263, 0.6596],
        [0.8465, 0.0030, 0.9253, 0.3342],
        [0.7469, 0.4648, 0.8200, 0.4193]])
>>> B=torch.split(A, 2, 0)
>>> B
(tensor([[0.3480, 0.4864, 0.2071, 0.6849],
        [0.5214, 0.6468, 0.9263, 0.6596]]), tensor([[0.8465, 0.0030, 0.9253, 0.3342],
        [0.7469, 0.4648, 0.8200, 0.4193]]))

在这个例子里，我们看到，原来 4x4 大小的 Tensor A，沿着第 0 维度，也就是沿“行”的方向，按照每组 2“行”的大小进行切分，得到了两个 2x4 大小的 Tensor。

那么问题来了，如果 split_size_or_sections 不能整除对应方向的大小的话，会有怎样的结果呢？我们将代码稍作修改就好了：

>>> C=torch.split(A, 3, 0)
>>> C
(tensor([[0.3480, 0.4864, 0.2071, 0.6849],
        [0.5214, 0.6468, 0.9263, 0.6596],
        [0.8465, 0.0030, 0.9253, 0.3342]]), tensor([[0.7469, 0.4648, 0.8200, 0.4193]]))

根据刚才的代码我们就能发现，原来，PyTorch 会尽可能凑够每一个结果，使得其对应 dim 的数据大小等于 split_size_or_sections。如果最后剩下的不够，那就把剩下的内容放到一块，作为最后一个结果。

接下来，我们再看一下 split_size_or_sections 是列表时的情况。刚才提到了，当 split_size_or_sections 为列表的时候，表示将此 tensor 切成和列表中元素大小一样的大小的块，我们来看一段对应的代码：

>>> A=torch.rand(5,4)
>>> A
tensor([[0.0871, 0.9387, 0.2978, 0.8540],
        [0.4216, 0.5009, 0.6090, 0.1782],
        [0.7486, 0.6665, 0.3248, 0.9010],
        [0.0457, 0.1507, 0.5208, 0.3595],
        [0.4709, 0.9482, 0.0524, 0.0906]])
>>> B=torch.split(A,(2,3),0)
>>> B
(tensor([[0.0871, 0.9387, 0.2978, 0.8540],
        [0.4216, 0.5009, 0.6090, 0.1782]]), tensor([[0.7486, 0.6665, 0.3248, 0.9010],
        [0.0457, 0.1507, 0.5208, 0.3595],
        [0.4709, 0.9482, 0.0524, 0.0906]]))

这部分代码怎么解释呢？其实也很好理解，就是将 Tensor A，沿着第 0 维进行切分，每一个结果对应维度上的尺寸或者说大小，分别是 2（行），3（行）。

unbind

通过学习前面的几个函数，咱们知道了怎么按固定大小做切分，或者按照索引 index 来进行选择。现在我们想象一个应用场景，如果我们现在有一个 3 channel 图像的 Tensor，想要逐个获取每个 channel 的数据，该怎么做呢？

假如用 chunk 的话，我们需要将 chunks 设为 3；如果用 split 的话，需要将 split_size_or_sections 设为 1。

虽然它们都可以实现相同的目的，但是如果 channel 数量很大，逐个去取也比较折腾。这时候，就需要用到另一个函数：unbind，它的函数定义如下：

torch.unbind(input, dim=0)

其中，input 表示待处理的 Tensor，dim 还是跟前面的函数一样，表示切片的方向。

我们结合例子来理解：

>>> A=torch.arange(16).view(4,4)
>>> A
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
>>> b=torch.unbind(A, 0)
>>>
>>> b
(tensor([0, 1, 2, 3]), tensor([4, 5, 6, 7]), tensor([ 8,  9, 10, 11]), tensor([12, 13, 14, 15]))

在这个例子中，我们首先创建了一个 4x4 的二维矩阵 Tensor，随后我们从第 0 维，也就是“行”的方向进行切分，因为矩阵有 4 行，所以就会得到 4 个结果。

接下来，我们看一下：如果从第 1 维，也就是“列”的方向进行切分，会是怎样的结果呢：

>>> b=torch.unbind(A, 1)
>>> b
(tensor([ 0,  4,  8, 12]), tensor([ 1,  5,  9, 13]), tensor([ 2,  6, 10, 14]), tensor([ 3,  7, 11, 15]))

不难发现，这里是按照“列”的方向进行拆解的。所以，unbind 是一种降维切分的方式，相当于删除一个维度之后的结果。

六、Tensor 的索引操作

你有没有发现，刚才我们讲的 chunk 和 split 操作，我们都是将数据整体进行切分，并获得全部结果。但有的时候，我们只需要其中的一部分，这要怎么做呢？一个很自然的想法就是，直接告诉 Tensor 我想要哪些部分，这种方法我们称为索引操作。

索引操作有很多方式，有提供好现成 API 的，也有用户自行定制的操作，其中最常用的两个操作就是 index_select 和 masked_select，我们分别去看看用法。

index_select

这里就需要 index_select 这个函数了，其定义如下：

torch.index_select(tensor, dim, index)

这里的 tensor、dim 跟前面函数里的一样，不再赘述。我们重点看一看 index，它表示从 dim 维度中的哪些位置选择数据，这里需要注意，index是 torch.Tensor 类型。

还是跟之前一样，我们来看几个示例代码：

>>> A=torch.arange(16).view(4,4)
>>> A
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])
>>> B=torch.index_select(A,0,torch.tensor([1,3]))
>>> B
tensor([[ 4,  5,  6,  7],
        [12, 13, 14, 15]])
>>> C=torch.index_select(A,1,torch.tensor([0,3]))
>>> C
tensor([[ 0,  3],
        [ 4,  7],
        [ 8, 11],
        [12, 15]])

在这个例子中，我们先创建了一个 4x4 大小的矩阵 Tensor A。然后，我们从第 0 维选择第 1（行）和 3（行）的数据，并得到了最终的 Tensor B，其大小为 2x4。随后我们从 Tensor A 中选择第 0（列）和 3（列）的数据，得到了最终的 Tensor C，其大小为 4x2。

怎么样，是不是非常简单？

masked_select

刚才介绍的 indexed_select，它是基于给定的索引来进行数据提取的。但有的时候，我们还想通过一些判断条件来进行选择，比如提取深度学习网络中某一层中数值大于 0 的参数。

这时候，就需要用到 PyTorch 提供的 masked_select 函数了，我们先来看它的定义：

torch.masked_select(input, mask, out=None)

这里我们只需要关心前两个参数，input 和 mask。

input 表示待处理的 Tensor。mask 代表掩码张量，也就是满足条件的特征掩码。这里你需要注意的是，mask 须跟 input 张量有相同数量的元素数目，但形状或维度不需要相同。

你是不是还感觉有些云里雾里？让我来举一个例子，你看了之后，一下子就能明白。

你在平时的练习中有没有想过，如果我们让 Tensor 和数字做比较，会有什么样的结果？比如后面这段代码，我们随机生成一个 5 位长度的 Tensor A：

>>> A=torch.rand(5)
>>> A
tensor([0.7484, 0.7311, 0.6890, 0.0034, 0.3469])
>>> B=A>0.3
>>> B
tensor([ True,  True,  True, False,  True])

在这段代码里，我们让 A 跟 0.3 做比较，得到了一个新的 Tensor，内部每一个数值表示的是 A 中对应数值是否大于 0.3。

比如第一个数值原来是 0.3731，大于 0.3，所以是 True；最后一个数值 0.2285 小于 0.3，所以是 False。

这个新的 Tensor 其实就是一个掩码张量，它的每一位表示了一个判断条件是否成立的结果。

然后，我们继续写一段代码，看看基于掩码 B 的选择是怎样的结果：

>>> C=torch.masked_select(A, B)
>>> C
tensor([0.7484, 0.7311, 0.6890, 0.3469])

你会发现，C 实际上得到的就是：A 中“满足 B 里面元素值为 True 的”对应位置的数据。

好了，这下你应该知道了 masked_select 的作用了吧？其实就是我们根据要筛选的条件，得到一个掩码张量，然后用这个张量去提取 Tensor 中的数据。

根据这个思路，上面的例子就可以简化为：

>>> A=torch.rand(5)
>>> A
tensor([0.8546, 0.7931, 0.5147, 0.7306, 0.2706])
>>> C=torch.masked_select(A, A>0.3)
>>> C
tensor([0.8546, 0.7931, 0.5147, 0.7306])

是不是非常简单呢？

再来做一个小练习：

现在有个 Tensor，如下：

>>> A=torch.tensor([[4,5,7], [3,9,8],[2,3,4]])
>>> A
tensor([[4, 5, 7],
        [3, 9, 8],
        [2, 3, 4]])

我们想提取出其中第一行的第一个，第二行的第一、第二个，第三行的最后一个，该怎么做呢？

很明显，通过masked_select()就可以做到：

>>> B = torch.tensor([[1,0,0],[1,1,0],[0,0,1]])
>>> C = B==1
>>> C
tensor([[ True, False, False],
        [ True,  True, False],
        [False, False,  True]])
>>> D = torch.masked_select(A, C)
>>> D
tensor([4, 3, 9, 4])

Tensor 中的主要函数跟用法还有很多，总结如下表，我们在使用的时候，根据需要灵活查询相关的参数列表即可：