【新手向】PyTorch常用Tensor shape变换方法

【新手向】PyTorch常用Tensor shape变换方法

前言

B站UP主科研水神大队长的视频中介绍了“缝合模块”大法，其中专门强调了“深度学习 玩的就是shape”。受此启发，专门整理能够调整tensor形状的几个内置函数，方便以后更好地调整PyTorch代码中的模型结构。

squeeze, unsqueeze

• torch.squeeze()
• torch.unsqueeze()

squeeze()用于压缩掉指定的维度，这个维度的取值必须是1，否则无效。unsqueeze()用于在指定的位置增加一个维度。

代码实例：

import torch

# image
x1 = torch.ones([4, 3, 256, 256]) # batch_size, channels, height, width
print("x1.shape:", x1.shape) # x1.shape: torch.Size([4, 3, 256, 256])

x2 = torch.ones([1, 1, 3])
print("x2.shape:", x2.shape) # x2.shape: torch.Size([1, 1, 3])

y1 = x1.squeeze(0)
print("y1.shape:", y1.shape) # y1.shape: torch.Size([4, 3, 256, 256])

y2 = x2.squeeze(0)
print("y2.shape:", y2.shape) # y2.shape: torch.Size([1, 3])

y6 = x2.squeeze(1)
print("y6.shape:", y6.shape) # y6.shape: torch.Size([1, 3])

y7 = x2.squeeze(-1)
print("y7.shape:", y7.shape) # y7.shape: torch.Size([1, 1, 3])

y3 = x1.unsqueeze(0)
print("y3.shape:", y3.shape) # y3.shape: torch.Size([1, 4, 3, 256, 256])

y4 = x1.unsqueeze(1)
print("y4.shape:", y4.shape) # y4.shape: torch.Size([4, 1, 3, 256, 256])

y5 = x1.unsqueeze(-1)
print("y5.shape:", y5.shape) # y5.shape: torch.Size([4, 3, 256, 256, 1])

transpose

transpose()用于调整tensor的维度顺序，在计算机视觉的任务中经常需要调整通道顺序，比如有的模型输出的顺序是(channel, height, width)，而有的输出顺序是(height, width, channel)，需要通过调换顺序来匹配输入输出。

transpose()有两种用法：

• torch.transpose()
• x.transpose()

代码实例：

import torch

x1 = torch.ones([4, 3, 256, 256]) # batch_size, channels, height, width
print("x1.shape:", x1.shape) # x1.shape: torch.Size([4, 3, 256, 256])

x2 = torch.ones([1, 1, 3])
print("x2.shape:", x2.shape) # x2.shape: torch.Size([1, 1, 3])

trans1 = torch.transpose(x1, 0, 1)
print("trans1.shape:", trans1.shape) # trans1.shape: torch.Size([3, 4, 256, 256])

trans2 = torch.transpose(x2, 1, 2)
print("trans2.shape:", trans2.shape) # trans2.shape: torch.Size([1, 3, 1])

trans3 = x1.transpose(0, 1)
print("trans3.shape:", trans3.shape) # trans3.shape: torch.Size([3, 4, 256, 256])

reshape

reshape()能够在总元素数量不产生变化的前提下改变tensor的形状。它也可以用于处理numpy array的形状。

代码实例：

import torch
import numpy as np

x3 = torch.Tensor([1, 2, 3, 4, 5, 6])
print("x3.shape:", x3.shape) # x3.shape: torch.Size([6])

reshape1 = x3.reshape(2, 3)
print("reshape1.shape:", reshape1.shape) # reshape1.shape: torch.Size([2, 3])

x4 = torch.ones([4, 4, 3, 256, 256])
print("x4.shape:", x4.shape) # x4.shape: torch.Size([4, 4, 3, 256, 256])

reshape2 = x4.reshape(4*4, 3, 256, 256)
print("reshape2.shape:", reshape2.shape) # reshape2.shape: torch.Size([16, 3, 256, 256])

x5 = np.array([1, 2, 3, 4, 5, 6])
print("x5.shape:", x5.shape) # x5.shape: (6,)

reshape3 = x5.reshape(2, 3)
print("reshape3.shape:", reshape3.shape) # reshape3.shape: (2, 3)

view

view()的作用与reshape()的作用相似，也是在总元素数量不产生变化的前提下改变形状，但view()只能对张量进行操作。

代码实例：

import torch

x3 = torch.Tensor([1, 2, 3, 4, 5, 6])
print("x3.shape:", x3.shape) # x3.shape: torch.Size([6])

view1 = x3.view(2, 3)
print("view1.shape:", view1.shape) # view1.shape: torch.Size([2, 3])
print("view1:", view1)
# view1: tensor([[1., 2., 3.],
#         [4., 5., 6.]])

view2 = x3.view(3, 2)
print("view2.shape:", view2.shape) # view2.shape: torch.Size([3, 2])
print("view2:", view2) 
# view2: tensor([[1., 2.],
#         [3., 4.],
#         [5., 6.]])

permute

permute()用于调整维度的顺序。与transpose()一次仅能“对调”两个维度的顺序不同，permute()可以一次调整多个维度的顺序。

代码实例：

import torch

x4 = torch.ones([4, 4, 3, 256, 256])
print("x4.shape:", x4.shape) # x4.shape: torch.Size([4, 4, 3, 256, 256])

permute1 = x4.permute(1, 3, 4, 0, 2)
print("permute1.shape:", permute1.shape) # permute1.shape: torch.Size([4, 256, 256, 4, 3])