PyTorch中的unsqueeze函数详解：给张量增加维度
什么是unsqueeze函数？
unsqueeze()函数的作用是在张量的指定位置插入一个维度，这个新维度的长度为1。
基本用法
1. 基本语法
import torch # 创建一个一维张量 x = torch.tensor([1, 2, 3, 4]) print(f"原始张量: {x}, 形状: {x.shape}") # torch.Size([4]) # 在维度0处增加维度 x_unsqueeze_0 = torch.unsqueeze(x, dim=0) print(f"在dim=0处unsqueeze: {x_unsqueeze_0}, 形状: {x_unsqueeze_0.shape}") # torch.Size([1, 4]) # 在维度1处增加维度（在最后一维之后） x_unsqueeze_1 = x.unsqueeze(dim=1) print(f"在dim=1处unsqueeze: {x_unsqueeze_1}, 形状: {x_unsqueeze_1.shape}") # torch.Size([4, 1])运行结果：
原始张量: tensor([1, 2, 3, 4]), 形状: torch.Size([4])
在dim=0处unsqueeze: tensor([[1, 2, 3, 4]]), 形状: torch.Size([1, 4])
在dim=1处unsqueeze: tensor([[1],
[2],
[3],
[4]]), 形状: torch.Size([4, 1])
2. 不同维度的张量示例
# 二维张量 x_2d = torch.tensor([[1, 2], [3, 4]]) print(f"原始2D张量形状: {x_2d.shape}") # torch.Size([2, 2]) # 在维度0处增加维度 result_0 = x_2d.unsqueeze(0) print(f"dim=0: {result_0.shape}") # torch.Size([1, 2, 2]) # 在维度1处增加维度 result_1 = x_2d.unsqueeze(1) print(f"dim=1: {result_1.shape}") # torch.Size([2, 1, 2]) # 在维度2处增加维度 result_2 = x_2d.unsqueeze(2) print(f"dim=2: {result_2.shape}") # torch.Size([2, 2, 1]) # 使用负索引（从后往前计数） result_neg = x_2d.unsqueeze(-1) print(f"dim=-1: {result_neg.shape}") # torch.Size([2, 2, 1])实际应用场景
1. 神经网络输入处理
神经网络通常需要特定维度的输入：
# 假设我们有一个批量大小为1的RGB图像 image = torch.randn(3, 32, 32) # [channels, height, width] # 卷积层期望的输入维度: [batch_size, channels, height, width] # 需要添加batch维度 image_with_batch = image.unsqueeze(0) # 形状: [1, 3, 32, 32] print(f"神经网络输入形状: {image_with_batch.shape}") # 批量处理多个图像 batch_size = 8 images = torch.randn(batch_size, 3, 32, 32) # 正确的输入形状