1. 从一次广播错误说起：为什么我们需要关心维度

那天下午，我正在调试一个图像分类模型，想把一个形状为[32, 1, 224, 224]的批次图像张量和一个形状为[32, 224, 224]的掩码张量做逐元素相乘。直觉上，这俩张量除了中间那个“1”，其他维度都对齐了，应该能直接运算。我信心满满地写下了images * masks，结果PyTorch毫不留情地抛出了一个错误：

RuntimeError: The size of tensor a (1) must match the size of tensor b (224) at non-singleton dimension 2

错误信息里的non-singleton dimension让我愣了一下。我仔细看了看，问题就出在那个“1”上。在PyTorch的广播机制里，维度为1的维度（也叫单例维度或“虚”维度）是特殊的，它可以自动扩展来匹配另一个张量对应维度的大小。但这里，我的masks张量在第二个维度（索引为1的位置）压根就没有维度，它只有三个维度[32, 224, 224]。PyTorch试图将images的第二个维度（大小为1）与masks的第二个维度（大小为224）进行广播对齐，但规则是，要么维度大小相等，要么其中一个为1。可现在是masks在那个位置“没有维度”，这就不符合广播规则了。

问题的根源在于维度不匹配。images是四维的[批次, 通道, 高, 宽]，而masks被我处理成了三维的[批次, 高, 宽]。要让它们相乘，我需要给masks在“通道”那个位置也插入一个维度，变成[32, 1, 224, 224]，这样广播机制就能愉快地工作了：[32, 1, 224, 224] * [32, 1, 224, 224]，其中masks的“1”会自动复制32次来匹配images的32个通道（虽然这里每个样本只有一个通道的掩码）。

这个看似简单的“插入一个维度”的操作，就是torch.unsqueeze()的用武之地。而它的逆操作，torch.squeeze()，则专门用来删除那些多余的、大小为1的维度，让张量变得更“紧凑”，避免在一些函数（如某些损失函数或矩阵运算）中因为多余的维度而报错。

unsqueeze和squeeze是PyTorch张量操作中最基础、最高频的两个函数。它们不改变张量的实际数据，只改变其“视图”，即我们对数据组织方式的解释。理解它们，是理解张量形状操作、广播机制乃至整个模型数据流的关键第一步。无论你是正在处理多模态数据（如图像+文本），需要对齐不同来源张量的维度；还是在搭建神经网络层，需要调整输入输出的形状；亦或是在进行简单的数据预处理，这两个函数都是你工具箱里的必备利器。

2. 维度的“增”与“删”：unsqueeze与squeeze的核心原理

要玩转unsqueeze和squeeze，首先得抛开对维度顺序的刻板印象。我们习惯用[batch, channel, height, width]来描述图像，用[batch, sequence_length, feature_dim]来描述序列。但维度本身只是一个索引数据的坐标轴。unsqueeze和squeeze操作的就是这些坐标轴。

2.1torch.unsqueeze(dim)：在指定位置插入一个新维度

unsqueeze的功能是在张量的指定维度索引dim处，插入一个大小为1的新维度。这里的dim指的是插入后新维度所在的位置。

关键理解：dim参数可以是负数。这是很多初学者困惑的地方。正数索引从左边开始（0-based），而负数索引从右边开始。dim=-1表示在最后一个维度之后插入，dim=-2表示在倒数第二个维度之前插入，以此类推。这个特性在你不确定张量总维数时非常有用。

举个例子，假设我们有一个三维张量t，形状为[2, 3, 4]（可以想象成2个样本，每个样本是3x4的矩阵）。

import torch t = torch.randn(2, 3, 4) print(t.shape) # torch.Size([2, 3, 4]) # 在维度0（最前面）插入，变成 [1, 2, 3, 4] t1 = t.unsqueeze(0) print(t1.shape) # torch.Size([1, 2, 3, 4]) # 在维度1（原维度0和1之间）插入，变成 [2, 1, 3, 4] t2 = t.unsqueeze(1) print(t2.shape) # torch.Size([2, 1, 3, 4]) # 在最后一个维度之后插入，变成 [2, 3, 4, 1] t3 = t.unsqueeze(-1) print(t3.shape) # torch.Size([2, 3, 4, 1]) # 在倒数第二个维度之前插入（即原最后一个维度之前），变成 [2, 3, 1, 4] t4 = t.unsqueeze(-2) print(t4.shape) # torch.Size([2, 3, 1, 4])

一个常见的应用场景：为全连接层准备数据。全连接层 (nn.Linear) 期望的输入是[batch_size, feature_size]。如果你有一个单独的样本，形状是[feature_size]（即一个一维向量），直接输入会报错，因为PyTorch默认第一个维度是batch。这时你需要unsqueeze(0)将其变为[1, feature_size]，表示批次大小为1。

2.2torch.squeeze(dim=None)：删除大小为1的维度

squeeze是unsqueeze的逆操作，它删除张量中所有大小为1的维度。如果指定了dim参数，则只尝试删除该特定维度，且仅当该维度大小为1时才生效，否则张量保持不变。

# 接上例，t1的形状是 [1, 2, 3, 4] print(t1.shape) # torch.Size([1, 2, 3, 4]) # 不指定dim，删除所有大小为1的维度，变回 [2, 3, 4] t1_squeezed = t1.squeeze() print(t1_squeezed.shape) # torch.Size([2, 3, 4]) # 注意：t1_squeezed 和最初的 t 在数据上是相同的（共享内存或经过复制）。 # 指定删除维度0（大小为1），效果同上 t1_squeezed_dim0 = t1.squeeze(0) print(t1_squeezed_dim0.shape) # torch.Size([2, 3, 4]) # 指定删除维度1（大小为2，不是1），所以张量不变 t1_squeezed_dim1 = t1.squeeze(1) print(t1_squeezed_dim1.shape) # torch.Size([1, 2, 3, 4]) # 形状未变 # 对于 t3 ([2, 3, 4, 1])，不指定dim会删除最后一个维度 t3_squeezed = t3.squeeze() print(t3_squeezed.shape) # torch.Size([2, 3, 4])

为什么需要squeeze？主要有两个原因：

1. 减少干扰：某些操作（如损失函数nn.CrossEntropyLoss）对输入形状有严格要求。例如，分类任务中，模型输出可能是[batch, num_classes, 1, 1]（在某些CNN结构后），而损失函数期望[batch, num_classes]，这时就需要squeeze掉后面两个为1的维度。
2. 节省内存与提升可读性：虽然大小为1的维度不增加数据量，但它们会在代码中传递，让张量的逻辑形状变得复杂。squeeze可以让张量形状更清晰，也避免在一些检查中引发意外。

注意：unsqueeze和squeeze返回的通常是原张量的一个视图（view），这意味着它们与原始张量共享底层数据存储，修改其中一个会影响另一个。这是一种高效的内存操作。但需要注意的是，如果squeeze或unsqueeze操作导致张量在内存中的连续性（contiguity）被破坏，某些后续操作（如view()）可能会要求你先调用.contiguous()。

2.3 原位操作与函数式操作

和大多数PyTorch张量操作一样，unsqueeze和squeeze都有两种使用方式：

• 函数式操作：torch.unsqueeze(input, dim)和torch.squeeze(input, dim=None)，返回一个新的张量。
• 原位操作：tensor.unsqueeze_(dim)和tensor.squeeze_(dim=None)，带下划线的版本会直接修改原张量。

x = torch.tensor([1, 2, 3]) y = x.unsqueeze(0) # y是新的张量，x不变 print(x.shape) # torch.Size([3]) print(y.shape) # torch.Size([1, 3]) x.unsqueeze_(0) # 直接修改x print(x.shape) # torch.Size([1, 3])

原位操作可以节省一点内存，但在计算梯度时需要小心，因为它会覆盖原变量的值。在神经网络的前向传播中，如果确定后续不再需要原张量，使用原位操作是安全的。但在需要保留计算图的情况下，建议使用函数式操作。

3. 实战场景深度剖析：从数据预处理到模型集成

理解了基本原理后，我们来看看unsqueeze和squeeze在真实项目中的用武之地。这些场景远比简单的维度加减要复杂和微妙。

3.1 场景一：图像与掩码的广播对齐（开篇问题的解决）

回到最初的问题。我们有图像张量images形状为[32, 1, 224, 224]，掩码张量masks形状为[32, 224, 224]。目标是让每个图像的每个通道（这里只有一个通道）与对应的掩码相乘。

错误做法：直接images * masks。因为维度不匹配，PyTorch无法广播。

正确做法：使用unsqueeze为masks添加一个通道维度。

# 假设 images: [32, 1, 224, 224], masks: [32, 224, 224] # 我们需要在masks的维度1（通道维）位置插入一个维度 masks_unsqueezed = masks.unsqueeze(1) # 形状变为 [32, 1, 224, 224] # 或者使用更清晰的写法，指明是通道维 # masks_unsqueezed = masks.unsqueeze(dim=1) result = images * masks_unsqueezed # 现在可以广播了 print(result.shape) # torch.Size([32, 1, 224, 224])

为什么是dim=1？因为在我们约定的图像张量形状[N, C, H, W]中，索引1的位置代表通道维。我们需要让masks在这个位置有一个维度（大小为1），以便与images的通道维（大小也为1）进行广播。

进阶思考：如果images是RGB三通道图[32, 3, 224, 224]，而masks仍然是单通道的[32, 224, 224]，我们仍然希望用同一个掩码作用于所有三个颜色通道。这时，masks.unsqueeze(1)得到[32, 1, 224, 224]，在与[32, 3, 224, 224]相乘时，masks在通道维上的1会自动扩展为3，实现“一对三”的掩码操作。这是广播机制的强大之处。

3.2 场景二：序列数据处理与注意力机制

在自然语言处理中，我们经常处理序列数据。假设我们有一批文本，经过嵌入层后得到张量embeddings，形状为[batch_size, seq_len, hidden_dim]，例如[16, 50, 768]。

现在我们要计算一个注意力权重向量attention_weights，它的形状是[batch_size, seq_len]，即[16, 50]，代表每个序列中每个词的重要性。我们想用这个权重对隐藏状态进行加权求和（通常称为注意力池化）。

一种简单的方法是：

# embeddings: [16, 50, 768] # attention_weights: [16, 50] # 我们需要将权重应用到最后一个维度（hidden_dim）上 # 第一步：将权重从 [16, 50] 变为 [16, 50, 1]，以便与embeddings广播 weights = attention_weights.unsqueeze(-1) # 形状 [16, 50, 1] # 第二步：逐元素相乘，权重会广播到768维 weighted_embeddings = embeddings * weights # 形状 [16, 50, 768] # 第三步：在序列长度维度上求和，得到加权后的句子表示 context_vector = weighted_embeddings.sum(dim=1) # 形状 [16, 768]

这里unsqueeze(-1)是关键一步，它在attention_weights的末尾添加了一个维度，使其形状变为[16, 50, 1]。这样，在与[16, 50, 768]相乘时，权重会在最后一个维度（隐藏维度）上自动复制768次，实现每个隐藏单元都按相同权重缩放的效果。

反过来，squeeze也经常出现在序列模型中。例如，一个双向LSTM最后一层的输出可能是[batch, seq_len, num_directions * hidden_size]，如果你只取最后一个时间步的输出 (output[:, -1, :])，它的形状是[batch, num_directions * hidden_size]。但如果你在处理序列分类任务时，用了nn.LSTM并设置了batch_first=True，且希望获取最后一个时间步的隐藏状态，你可能会得到形状为[batch, 1, hidden_size]的张量（取决于如何索引）。为了送入后续的全连接层，你需要squeeze(1)去掉中间的维度。

3.3 场景三：损失函数输入的形状适配

这是新手踩坑的重灾区。以交叉熵损失nn.CrossEntropyLoss为例，它要求两个输入：

• input：模型的原始输出（未经过Softmax），形状为[batch_size, num_classes]。
• target：真实标签，形状为[batch_size]，每个值是类别索引（0到num_classes-1）。

假设你有一个简单的CNN用于MNIST分类（10类），最后一层是nn.Linear(512, 10)。前向传播后，你得到的output形状是[batch, 10]，这很好。

但如果你用的网络结构比较复杂，比如在全局平均池化后，你可能会得到一个形状为[batch, 10, 1, 1]的输出（这在一些迁移学习模型中很常见）。

# 模拟一个“奇怪”的输出形状 output = torch.randn(32, 10, 1, 1) # [batch, classes, 1, 1] target = torch.randint(0, 10, (32,)) # [batch] loss_fn = nn.CrossEntropyLoss() # 直接计算会报错：4D target tensor is not supported # loss = loss_fn(output, target) # RuntimeError! # 正确做法：squeeze掉后面两个为1的维度 output_corrected = output.squeeze() # 形状变为 [32, 10] loss = loss_fn(output_corrected, target) # 现在可以了 print(loss)

同样的问题也出现在target上。如果你不小心把target做成了[batch, 1]的形状（例如从DataLoader中取出时没有处理），也需要squeeze()一下。

经验之谈：在将张量送入损失函数之前，养成检查形状的习惯。对于CrossEntropyLoss，一个简单的断言很有用：assert input.shape == (batch, num_classes) and target.shape == (batch,)。如果形状不对，squeeze和unsqueeze就是你快速修复的工具。

3.4 场景四：自定义层与维度兼容性

当你编写自定义的PyTorch层或函数时，考虑输入维度的灵活性是一个好习惯。你的函数可能期望某种特定维度的输入，但用户可能传入不同维度的数据。这时，unsqueeze/squeeze可以帮助你进行内部标准化。

例如，你写了一个计算逐元素高斯加权的函数，它期望输入是[..., H, W]的形式（...表示任意多的批次维度），并对最后两个空间维度进行加权。

def spatial_gaussian_weight(x): """ 对输入x的最后两个维度进行高斯加权。 x: 张量，形状为 [..., H, W] 返回: 加权后的张量，形状不变 """ # 假设我们有一个简单的二维高斯核，形状为 [H, W] H, W = x.shape[-2], x.shape[-1] # 这里简化创建高斯核的过程 gaussian_kernel = torch.randn(H, W) # 仅为示例，实际应为高斯分布 # 为了进行逐元素相乘，我们需要将核扩展到与x相同的维度 # 但x可能有前面的批次维度，如 [B, C, H, W] 或 [B, H, W] # 我们需要让gaussian_kernel的形状变为 [1, 1, H, W] 或 [1, H, W] 以便广播 # 一个通用的方法是：在核的前面添加足够的维度1，直到其维度和x一样多 while gaussian_kernel.dim() < x.dim(): gaussian_kernel = gaussian_kernel.unsqueeze(0) # 现在gaussian_kernel的形状比如是 [1, 1, H, W] (如果x是4D) # 或者 [1, H, W] (如果x是3D) return x * gaussian_kernel # 测试 x_4d = torch.randn(2, 3, 5, 5) # [B, C, H, W] x_3d = torch.randn(2, 5, 5) # [B, H, W] print(spatial_gaussian_weight(x_4d).shape) # torch.Size([2, 3, 5, 5]) print(spatial_gaussian_weight(x_3d).shape) # torch.Size([2, 5, 5])

在这个函数中，unsqueeze(0)被循环使用，动态地将二维高斯核的维度提升到与输入x相同，从而实现了灵活的广播。这使得函数能够处理不同批次和通道维度的输入，增强了代码的鲁棒性。

4. 高级技巧、常见陷阱与性能考量

掌握了基本应用后，我们来看看一些更深入的知识点和容易踩的坑。

4.1view()、reshape()与squeeze/unsqueeze的协同与区别

view()和reshape()也可以改变形状，但它们与squeeze/unsqueeze有本质区别：

• squeeze/unsqueeze：只操作大小为1的维度，进行纯粹的维度增减，不改变元素间的相对顺序和总数。它们是“安全”的形状变换，通常能返回一个视图。
• view()/reshape()：可以改变任意维度的大小，但必须保证变换前后元素总数一致。它们可能会改变数据在内存中的布局（view()要求张量是连续的，否则会报错；reshape()在可能的情况下返回视图，否则返回拷贝）。

它们经常结合使用：

# 将一个三维张量 [2, 3, 4] 扁平化为二维 [2, 12] x = torch.randn(2, 3, 4) # 方法1: 先用 view 展平后两维，但需要知道具体大小 x_flat1 = x.view(2, -1) # -1 表示自动推断 # 方法2: 更通用的方法，适用于不知道后几维总大小的情况 # 先合并后两维，再展平 x_squeezed = x.flatten(start_dim=1) # 从维度1开始展平，得到 [2, 12] # flatten 内部其实做了类似 view 的操作 # 一个更复杂的例子：将 [B, C, H, W] 转换为 [B, C*H*W] 以送入全连接层 batch, chan, height, width = x_4d.shape x_fc_input = x_4d.view(batch, -1) # 常用 # 或者 x_fc_input = x_4d.flatten(1) # 更清晰，从第1维（C）开始展平 # 转换回来后，可能需要 unsqueeze # 假设我们从全连接层得到一个 [B, C] 的输出，想把它变成 [B, C, 1, 1] 以模拟空间维度 fc_output = torch.randn(batch, chan) spatial_output = fc_output.view(batch, chan, 1, 1) # 这等价于： spatial_output = fc_output.unsqueeze(-1).unsqueeze(-1) # 或者 spatial_output = fc_output[:, :, None, None] # 使用None索引，是unsqueeze的语法糖

关键区别：view()和reshape()关心的是形状的重新排列，而squeeze和unsqueeze关心的是维度的存在与否。当你只是需要增加或删除一个“虚”维度时，用后者更直观、更安全。

4.2 使用None索引进行快速维度操作

PyTorch（和NumPy）支持使用None（在NumPy中是np.newaxis）在索引中插入新维度，这本质上是unsqueeze的语法糖，非常简洁。

x = torch.randn(2, 3) print(x.shape) # [2, 3] # 在维度0插入 y1 = x[None, :, :] # 等价于 x.unsqueeze(0) print(y1.shape) # [1, 2, 3] # 在维度1插入 y2 = x[:, None, :] # 等价于 x.unsqueeze(1) print(y2.shape) # [2, 1, 3] # 在最后一个维度之后插入 y3 = x[:, :, None] # 等价于 x.unsqueeze(-1) print(y3.shape) # [2, 3, 1] # 甚至可以同时插入多个 y4 = x[None, :, :, None] # 等价于 x.unsqueeze(0).unsqueeze(-1) print(y4.shape) # [1, 2, 3, 1]

这种方式在临时需要增加维度进行广播计算时特别方便，代码更紧凑。但它的可读性略差，尤其是对于不熟悉该语法的读者。在重要的、需要清晰表达的代码中，我倾向于使用unsqueeze，因为函数名本身就是文档。

4.3 内存连续性（Contiguity）的潜在影响

这是一个高级但重要的主题。PyTorch张量在内存中的存储有“连续”和“不连续”之分。像transpose(),permute(),narrow(),select()等操作返回的是原张量的视图，它们改变了索引方式，但可能破坏了内存的连续性。而view()严格要求张量是连续的，否则会报错。reshape()会尝试返回视图，如果不行就返回拷贝。

squeeze和unsqueeze通常返回视图，而且它们一般不会破坏连续性，因为它们只是增加或删除一个大小为1的维度，不改变元素顺序。但是，如果你对一个不连续张量进行squeeze/unsqueeze操作，结果张量也可能是不连续的。

在绝大多数情况下，你不需要关心这个。但如果你在squeeze/unsqueeze之后立即调用view()，或者进行一些需要连续内存的低级操作（如与C/C++扩展交互），可能会遇到问题。

x = torch.randn(2, 3, 4) y = x.transpose(1, 2) # y的形状是[2, 4, 3]，并且是不连续的 print(y.is_contiguous()) # False z = y.unsqueeze(1) # z的形状是[2, 1, 4, 3] print(z.is_contiguous()) # 可能仍然是False # 如果此时想用 view 改变形状，可能会报错 # w = z.view(2, -1) # 可能报错: RuntimeError: view size is not compatible... # 安全的做法是先使其连续 w = z.contiguous().view(2, -1) print(w.shape) # torch.Size([2, 12])

经验法则：如果你进行了一系列复杂的维度变换操作（特别是包含transpose,permute），然后在最后需要改变形状时，先调用.contiguous()再view()是更稳妥的做法。对于squeeze/unsqueeze单独使用，则基本不用担心。

4.4 常见陷阱与调试技巧

1. dim参数索引错误：这是最常见的错误。时刻记住dim参数指的是操作后的维度索引位置。对于unsqueeze(dim)，dim的范围是[-input.dim()-1, input.dim()]。一个有用的调试方法是打印操作前后的shape。

   x = torch.randn(5, 10) print(f'Original shape: {x.shape}') # [5, 10] for dim in range(-x.dim()-1, x.dim()+1): try: y = x.unsqueeze(dim) print(f'unsqueeze(dim={dim:2d}) -> shape: {y.shape}') except IndexError as e: print(f'unsqueeze(dim={dim:2d}) -> Error: {e}')

2. 过度squeeze：使用无参数的squeeze()会删除所有大小为1的维度。这有时会过度删除你本想保留的维度。例如，一个形状为[1, 10, 1, 20]的张量，经过squeeze()会变成[10, 20]，丢失了批次和通道信息。更好的做法是明确指定dim参数，只删除你确定无用的维度。

3. 与size()和shape属性的混淆：x.size(dim)返回第dim维的大小，而x.dim()返回总维数。在动态计算unsqueeze的dim参数时，这些方法很有用。

   # 动态地在倒数第二个维度之前插入 x = torch.randn(3, 4, 5) insert_dim = -2 # 等价于 x.dim() - 1 y = x.unsqueeze(insert_dim) print(y.shape) # [3, 4, 1, 5]

4. 广播语义理解不清：unsqueeze的核心目的是为了广播。如果你unsqueeze后仍然无法进行运算，请仔细检查两个张量的形状，在纸上画出它们的维度，并回想广播规则：从后往前对齐维度，每个维度要么相等，要么其中一个为1，要么其中一个不存在（可以理解为1）。unsqueeze就是为了让“不存在”变成“存在且为1”。

5. 原地操作 (_方法) 的梯度：在自定义nn.Module的forward方法中，如果对需要求导的张量使用了squeeze_()或unsqueeze_()，这不会破坏梯度传播，因为这只是改变了张量的元数据（形状、步长等），而不是数据本身。但出于代码清晰和避免副作用的考虑，在大多数情况下使用非原位操作是更好的选择。

unsqueeze和squeeze就像张量维度世界的“精细手术刀”，虽小但不可或缺。它们不直接参与复杂的数学计算，却是构建正确数据流、连接不同模块的桥梁。从数据加载时的维度调整，到模型内部的特征变换，再到损失计算前的形状适配，几乎贯穿了深度学习项目的数据处理全流程。理解并熟练运用它们，能让你在调试“The size of tensor a must match the size of tensor b”这类错误时更加得心应手，写出更健壮、更清晰的PyTorch代码。下次当你需要对维度动手脚时，先问问自己：我是需要增加一个维度来广播，还是需要删除一个冗余的维度来简化？想清楚了这个问题，unsqueeze和squeeze自然会成为你的得力助手。