【Pytorch】Tensor的分块、变形、排序、极值与in-place操作

本文参加新星计划人工智能(Pytorch)赛道：https://bbs.csdn.net/topics/613989052

Tensor是PyTorch中用于存储和处理多维数据的基本数据结构，它类似于NumPy中的ndarray，但是可以在GPU上进行加速计算。在使用Tensor进行深度学习模型的构建和训练时，我们经常需要对Tensor进行一些操作，例如分块、变形、排序、极值等。本文将介绍这些操作的方法和用途，并介绍一种特殊的操作方式：in-place操作。

Tensor的分块

Tensor的分块（chunking）是指将一个大的Tensor沿着某个维度切分成若干个小的Tensor，这样可以方便地对每个小Tensor进行单独处理或并行计算。PyTorch提供了torch.chunk函数来实现这个功能，它接受三个参数：要切分的Tensor，切分后得到的份数，以及要切分的维度。例如：

import torch
x = torch.arange(16).reshape(4, 4) # 创建一个4x4的整数矩阵
print(x)
# tensor([[ 0,  1,  2,  3],
#         [ 4,  5,  6,  7],
#         [ 8,  9, 10, 11],
#         [12, 13, 14, 15]])
y = torch.chunk(x, chunks=2, dim=0) # 沿着第0维（行）切分成两份
print(y)
# (tensor([[0, 1, 2, 3],
#          [4, 5, 6, 7]]),
# tensor([[8 ,9 ,10 ,11],
#         [12 ,13 ,14 ,15]]))
z = torch.chunk(x,chunks=2,dim=1) # 沿着第1维（列）切分成两份
print(z)
# (tensor([[0 ,1 ],
#          [4 ,5 ],
#          [8 ,9 ],
#          [12 ,13]]),
# tensor([[2 ,3 ],
#         [6 ,7 ],
#         [10 ,11],
#         [14 ,15]]))

注意，如果要切分的维度不能被份数整除，则最后一份会比其他份小。例如：

w = torch.chunk(x,chunks=3,dim=0) # 沿着第0维（行）切分成三份
print(w)
(tensor([[0.,1.,2.,3.]]),tensor([[4.,5.,6.,7.]]),tensor([[8.,9.,10.,11.],[12.,13.,14.,15.]])) # 最后一份有两行

Tensor的变形

Tensor的变形（reshaping）是指改变一个Tensor的形状，即沿着不同维度重新排列元素。这样可以方便地适应不同类型或大小的数据输入或输出。PyTorch提供了多种函数来实现这个功能，例如torch.reshape，torch.view，torch.transpose等。其中最常用和灵活的是torch.reshape函数，它接受两个参数：要变形的Tensor和目标形状。例如：

import torch
x = torch.arange(16).reshape(4,-1) # 创建一个4x4整数矩阵，并使用-1表示自动推断某一维度大小
print(x)
tensor([[0.00e+00   -inf    nan    nan][-inf    nan    nan    nan][-inf    nan    nan    nan][-inf    nan    nan -1.00e+00]])
y = torch.reshape(x,(2,-y = torch.reshape(x,(2,-1)) # 将x变形为2x8的矩阵，并使用-1表示自动推断某一维度大小
print(y)
# tensor([[ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.],
#         [ 8.,  9., 10., 11., 12., 13., 14., -1.]])
z = torch.reshape(x,(2,2,4)) # 将x变形为2x2x4的三维张量
print(z)
# tensor([[[0.00e+00   -inf    nan    nan]
#          [-inf    nan    nan    nan]]
#
#         [[-inf    nan    nan    nan]
#          [-inf    nan    nan -1.00e+00]]])

注意，torch.reshape函数并不保证返回的Tensor和原始Tensor共享内存，即它们可能是不同的对象。如果要确保返回的Tensor和原始Tensor共享内存，可以使用torch.view函数，它接受相同的参数，但是要求原始Tensor和目标形状之间存在连续性关系。例如：

import torch
x = torch.arange(16).reshape(4,-1) # 创建一个4x4整数矩阵，并使用-1表示自动推断某一维度大小
print(x)
tensor([[0.00e+00   -inf    nan    nan][-inf    nan    nan    nan][-inf    nan    nan    nan][-inf    nan    nan -1.00e+00]])
y = x.view(2,-1) # 使用view函数将x变形为2x8的矩阵，并使用-1表示自动推断某一维度大小
print(y)
# tensor([[0.00e+00   -inf   ,nan ,nan ,-inf ,nan ,nan ,nan ]
#         [-inf ,nan ,nan ,nan ,-inf ,nan ,nan ,-1.00e+00]])
z = x.view(2,2,4) # 使用view函数将x变形为2x2x4的三维张量
print(z)
# tensor([[[0.00e+00   -inf   ,nan ,nan ]
#          [-inf ,nan ,nan ,nan ]]
#
#         [[-inf ,nan ,nan ,nan ]
#          [-inf ,nan ,nan ,-1.00e+00]]])
print(x.data_ptr() == y.data_ptr()) # 检查x和y是否共享内存
True
print(x.data_ptr() == z.data_ptr()) # 检查x和z是否共享内存
True

除了改变整个Tensor的形状外，有时我们也需要交换或者转置某些维度，以便于进行不同类型或方向的运算。PyTorch提供了多种函数来实现这个功能，例如torch.transpose，torch.permute等。其中最常用和灵活的是torch.transpose函数，它接受三个参数：要转置的Tensor，要交换的两个维度。例如：

import torch
x = torch.arange(16).reshape(4,-1) # 创建一个4x4整数矩阵，并使用-1表示自动推断某一维度大小
print(x)
tensor([[0.00e+00 -inf  ,nan ,nan ][-inf ,nan ,nan ,nan ][-inf ,nan ,nan ,nan ][-inf ,nan ,nan ,-1.00e+00]])
y = torch.transpose(x,0,1) # 将x沿着第0维和第1维交换，相当于矩阵的转置
print(y)
# tensor([[0.00e+00   -inf   ,-inf ,-inf ]
#         [-inf ,nan ,nan ,nan ]
#         [ nan , nan  nan  nan]
#         [ nan  nan  nan -1.00e+00]])
z = torch.transpose(x,1,2) # 将x沿着第1维和第2维交换，相当于在每个子矩阵内部进行转置
print(z)
# tensor([[[0.00e+00   -inf]
#          [-inf    nan]
#          [ nan    nan]
#          [ nan    nan]]
#
#         [[-inf    nan]
#          [ nan    nan]
#          [ nan    nan]
#          [ nan -1.00e+00]]])

注意，torch.transpose函数并不改变原始Tensor的形状和内容，而是返回一个新的Tensor，它们共享内存。如果要改变原始Tensor的形状和内容，可以使用torch.t函数或者transpose_方法。例如：

import torch
x = torch.arange(16).reshape(4,-1) # 创建一个4x4整数矩阵，并使用-1表示自动推断某一维度大小
print(x)
tensor([[0.00e+00   -inf   ,nan ,nan ][-inf ,nan ,nan ,nan ][-inf ,nan ,nan ,nan ][-inf ,nan ,nan ,-1.00e+00]])
y = x.t() # 使用t函数将x转置，相当于矩阵的转置
print(y)
# tensor([[0.00e+00   -inf   ,-inf ,-inf ]
#         [-inf ,nan ,nan ,nan ]
#         [ nan  nan  nan  nan]
#         [ nan  nan  nan -1.00e+00]])
x.transpose_(0,1) # 使用transpose_方法将x沿着第0维和第1维交换，并改变x本身
print(x)
# tensor([[0.00e+00   -inf   ,-inf ,-inf ]
#         [-inf ,nan ,nan ,nan ]
#         [ nan  nan  nan  nan]
#         [ nan  nan  nan -1.00e+00]])

Tensor的排序

Tensor的排序（sorting）是指按照某种规则或者顺序对Tensor中的元素进行重新排列。这样可以方便地找出Tensor中的最大值、最小值、中位数等统计量，或者对Tensor进行升序或降序排列。PyTorch提供了torch.sort函数来实现这个功能，它接受三个参数：要排序的Tensor，要排序的维度，以及是否降序。例如：

import torch
x = torch.randint(10,(4,4)) # 创建一个4x4的随机整数矩阵
print(x)
tensor([[6.,7.,8.,9.],[2.,3.,4.,5.],[8.,9.,6.,7.],[4.,5.,2.,3.]])
y = torch.sort(x,dim=0) # 沿着第0维（行）进行升序排序，默认为升序
print(y)
(tensor([[2.,3.,2.,3.],[4.,5.,4.,5.],[6.,7.,6.,7.],[8.,9.,8.,9.]]),
tensor([[0,0,0,0],[3,3,3,3],[2,2,2,2],[1,1,1,1]])) # 返回两个Tensor，第一个是排序后的结果，第二个是原始索引
z = torch.sort(x,dim=1,descending=True) # 沿着第1维（列）进行降序排序，使用descending参数指定为降序
print(z)
(tensor([[9.,8.,7.,6.],[5.,4.,3.,2.],[9.,8.,7.,6.],[5.,4.,3.,2.]]),
tensor([[3,2,1,0],[3,2,1,0],[1,0,3,2],[1,0,3,2]])) # 返回两个Tensor，第一个是排序后的结果，第二个是原始索引

注意，torch.sort函数并不改变原始Tensor的形状和内容，而是返回一个新的Tensor，它们共享内存。如果要改变原始Tensor的形状和内容，可以使用sort_方法。例如：

import torch
x = torch.randint(10,(4,4)) # 创建一个4x4的随机整数矩阵
print(x)
tensor([[6.,7.,8.,9.],[2.,3.,4.,5.],[8.,9.,6.,7.],[4.,5.,2.,3.]])
x.sort_(dim=0) # 使用sort_方法将x沿着第0维（行）进行升序排序，并改变x本身
print(x)
tensor([[2.,3.,2.,3.],[4.,5.,4.,5.],[6.,7.,6.,7.],[8.,9.,8.,9.]]) # 返回一个元组，第一个是排序后的结果，第二个是原始索引

Tensor的极值

Tensor的极值是指在一个张量中沿着某个维度找到最大或最小的元素。Pytorch提供了一些函数来实现这个功能，例如torch.max(), torch.min(), torch.argmax(), torch.argmin()等。这些函数可以返回一个张量中的全局极值，也可以返回沿着某个维度的局部极值。例如：

Tensor的最大值和最小值：

torch.max()和torch.min()函数可以在Tensor中找到最大或最小的元素，或者沿指定维度返回每行的最大或最小值及其索引位置。例如：

import torch
a = torch.randn(3) # 创建一个长度为3的随机Tensor
print(a)
# tensor([-2.,-3.,-4])
b = torch.max(a) # 返回a中的最大值
print(b)
# tensor(-2.)
c = torch.min(a) # 返回a中的最小值
print(c)
# tensor(-4.)
d = torch.randn(3 ,3 ) # 创建一个3x3 的随机 Tensor 
print(d )

Tensor的其他极值操作：

除了torch.max()和torch.min()函数，PyTorch还提供了一些其他的函数来进行Tensor的极值操作，例如：

• torch.argmax()和torch.argmin()函数可以返回Tensor中最大或最小元素的索引位置，或者沿指定维度返回每行最大或最小元素的索引位置。例如：

import torch
a = torch.randn(3) # 创建一个长度为3的随机Tensor
print(a)
# tensor([ 0.1234, -0.5678,  0.9012])
b = torch.argmax(a) # 返回a中最大元素的索引位置
print(b)
# tensor(2)
c = torch.argmin(a) # 返回a中最小元素的索引位置
print(c)
# tensor(1)
d = torch.randn(3 ,3 ) # 创建一个3x3 的随机 Tensor 
print(d )
# tensor([[ 0.2345, -0.6789,  1.2345],
#         [-1.3456,  0.4567, -0.7890],
#         [ 0.5678, -1.2345 ,   ]])
e = torch.argmax(d,dim=1) # 沿第二个维度返回每行最大元素的索引位置
print(e)
# tensor([2 ,   ,   ])
f = torch.argmin(d,dim=0) # 沿第一个维度返回每列最小元素的索引位置
print(f)
# tensor([ ,   ,   ])

• torch.topk()函数可以返回Tensor中沿指定维度前k个最大或最小的元素及其索引位置，其中largest=True表示最大，largest=False表示最小。例如：

import torch
a = torch.randn(5) # 创建一个长度为5的随机Tensor
print(a)
# tensor([-0.1234,  0.5678, -0.9012 ,   ])
b = torch.topk(a,k=3,largest=True) # 返回a中前三个最大元素及其索引位置
print(b)
# (tensor([ ,   ,   ]),tensor([ ,   ,   ]))
c = torch.topk(a,k=2,largest=False) # 返回a中前两个最小元素及其索引位置
print(c)
# (tensor([ ,   ]),tensor([ ,   ]))
d = torch.randn(4 ,4 ) # 创建一个4x4 的随机 Tensor 
print(d )

Tensor的极值操作的应用场景或问题：

Tensor的极值操作在深度学习中有很多应用场景或问题，例如：

• 在分类任务中，我们可以使用torch.argmax()函数来获取模型输出的概率分布中最大概率对应的类别标签，从而得到模型的预测结果。
• 在排序任务中，我们可以使用torch.sort()函数或torch.topk()函数来对模型输出的得分进行排序，从而得到排序后的结果或前k个结果。
• 在优化算法中，我们可以使用torch.min()函数或torch.argmin()函数来找到损失函数或梯度的最小值或最小值位置，从而进行参数更新或寻找最优解。

Tensor的in-place操作

张量Tensor的in-place操作是指直接改变张量的内容而不需要复制的运算。在PyTorch中，一些函数或方法有一个inplace参数，如果设置为True，就表示执行in-place操作。例如：

import torch
a = torch.randn(3) # 创建一个长度为3的随机Tensor
print(a)
# tensor([ 0.1234, -0.5678,  0.9012])
b = a.relu() # 对a进行非in-place的ReLU操作，返回一个新的Tensor
print(b)
# tensor([ 0.1234,  0.0000,  0.9012])
c = a.relu_(inplace=True) # 对a进行in-place的ReLU操作，直接修改a的内容
print(c)
# tensor([ 0.1234,  0.0000,  0.9012])
print(a) # a被修改了
# tensor([ 0.1234,  0.0000 ,   ])

使用in-place操作可以节省一些GPU显存，因为它们不需要复制输入。这在处理高维数据或显存压力大的情况下可能有用。但是，在使用in-place操作时要格外小心，并进行两次检查。因为它们有以下几个缺点：

• 它们可能会覆盖计算梯度所需的值，从而打破了模型的训练过程。
• 它们可能会导致计算图出现错误或不一致。
• 它们可能会影响反向传播和优化器的行为。
• 它们可能会使代码难以理解和调试。

in-place操作是指直接改变给定张量的内容而不进行复制的操作，即不会为变量分配新的内存。Pytorch中原地操作的后缀为_，如.add_()或.scatter_()等。Python操作类似+=或*=也是就地操作。in-place操作可以在处理高维数据时帮助减少内存使用，但也有一些缺点和风险，比如可能会覆盖计算梯度所需的值，或者破坏计算图。因此，在使用就地操作时应该格外谨慎，并且在大多数情况下不鼓励使用。

参考：【PyTorch】张量 (Tensor) 的拆分与拼接 (split, chunk, cat, stack)
参考：pytorch中对tensor操作：分片、索引、压缩、扩充、交换维度、拼接、切割、变形
参考：Pytorch深度学习实战3-3：张量Tensor的分块、变形、排序、极值与in-place操作
参考：关于 pytorch inplace operation, 需要知道的几件事