PyTorch的 torch.unsqueeze() 和 torch.squeeze()方法详解

torch.unsqueeze()

• 定义: torch.unsqueeze(input, dim) 返回一个新的张量，该张量为指定位置dim增加了一个大小为1的维度。

• 参数:

  • input (Tensor): 输入张量。
  • dim (int): 插入维度的索引。

• 返回值: Tensor。一个新的张量，它与输入张量input有着几乎相同的数据，但是增加了一个维度。

• 用途: 增加张量的一个维度，常在希望增加一个批处理维度时使用。

• 代码示例:

  import torch
  
  x = torch.tensor([1, 2, 3, 4])
  y = torch.unsqueeze(x, 0)  # 结果为：[[1, 2, 3, 4]]
  

  torch.squeeze()

• 定义: torch.squeeze(input, dim=None) 返回一个新的张量，它会移除input中所有指定维度dim的1维。

• 参数:

  • input (Tensor): 输入张量。
  • dim (int, 可选): 要压缩的维度。如果未指定，则移除所有单维度。

• 返回值: Tensor。一个新的张量，它与输入张量input有着相同的数据，但是移除了所有指定的单维度。

• 用途: 减少张量的维度，常用于移除批处理维度或任何大小为1的维度。

• 代码示例:

  import torch
  
  x = torch.tensor([[1, 2, 3, 4]])
  y = torch.squeeze(x, 0)  # 结果为：[1, 2, 3, 4]
  

       注意：在实际代码中，返回的张量通常会共享输入张量的数据，除非必须复制（比如要求返回一个不同内存的张量）。