说到深度学习的张量操作,PyTorch、TensorFlow、NumPy 里最容易让人犯迷糊的,就是 Stack 和 Concat。我在写 PyTorch 系列笔记第 28 篇时,正好帮 A 同学排查一个数据加载器的报错,他把四张 [3,224,224] 的图直接 torch.cat 到一起,结果得到一张 [12,224,224] 的“大图”,模型输入维度怎么都不对。这个问题的根源,就是没有搞清 cat 和 stack 在维度上的本质区别。这篇内容适合刚接触深度学习、或者正在从 NumPy 往 PyTorch 迁移的读者,看完你就能分清这两个操作,遇到维度报错时也知道该从哪里排查。
1. 先把错误示范讲透:Batch 为什么不能靠 Concat 拼出来
1.1 “拼接图片”为什么会得到错误的形状
我们经常要从一组图片构建一个 batch。假设有 4 张图片,每张读进来是 [3,224,224],也就是通道数、高、宽三个维度。看起来“把 4 张图片拼起来”最自然的写法是用 torch.cat(imgs, dim=0)。但错在哪里?我们直接跑一下:
import torch imgs = [torch.randn(3, 224, 224) for _ in range(4)] # 错误写法:得到 [12, 224, 224] wrong = torch.cat(imgs, dim=0) print(wrong.shape) # 正确写法:得到 [4, 3, 224, 224] right = torch.stack(imgs, dim=0) print(right.shape)cat 是在已有维度上做合并。这些图片没有 batch 维,dim=0 实际上是通道维,所以 cat 之后通道数变成了 3×4=12。站在图像语义上,这不是构建了一个 batch,而是把 4 张图的通道全部并到了一起,相当于一张有 12 个通道的“拼接图”。而 stack 会先在 0 位插入一个大小为 4 的新维度,再把每张图原本的 [3,224,224] 放进去,得到 [4,3,224,224]。这才是模型期望的 batch 形状。
反过来,如果你的数据已经是 [4,3,224,224] 这种 batch 形式,你想把两个 batch 合并成一个更大的 batch,才会用 cat(dim=0)。这也是很多训练代码里,把多次采样或不同进程送来的样本合并时用的方式。我们要区分“样本集合”和“batch 张量”:样本集合是没有 batch 维的列表,batch 张量是有 batch 维的数组,两者之间的转换是 stack,不是 cat。
1.2 模型输入维度对不上的排查习惯
模型前向传播报维度错误时,第一件事不是去 model forward 里找,而是检查 dataloader 输出。打印 batch 的 shape、dtype、device,通常很快就能定位。常见错误是:模型输入是 [N,C,H,W],而 dataloader 返回的却是 [N*C,H,W]。这大概率就是误用 cat 了。我用过一个简单规则:凡是“一堆单独的样本”要变成 batch,一定是用 stack;凡是“已经成型的 batch 要合并”,才用 concat。这个规则能避开绝大多数坑。
顺便提醒一句:在 TensorFlow 中默认布局是 [H,W,C],NumPy 读图也可能是 [H,W,C],但规则完全一样。一组 [224,224,3] 的图片要用 tf.stack 得到 [N,224,224,3],用 tf.concat 会得到 [224,224,3N],同样是错误的。这个问题不局限于某一个框架,理解了维度语义之后,三个库都通用。
1.3 为什么新手总在第一步栽跟头
因为“拼起来”这个直觉太强了。列表用 append 会增加一个元素,但维度不变;而深度学习中我们面对的是张量,还要区分“沿着已有维度延长”和“新增一个维度再堆放”。如果不建立“维度数量(秩)”这个意识,靠硬记 API 很容易记混。比如 NumPy 的 np.vstack 好像也能把一维数组堆成二维,但底层其实是先升维再 concat,这又让新手更难分清。后面我会专门拆开 stack 的内部过程,看完你就不会再被这些名字带偏。
2. Concat 和 Stack 的本质差异:一个在扩容,一个在加维
2.1 从“秩”的角度理解
“秩”可以理解为形状里有多少个数字。一个向量 [3] 是 1 维,一个矩阵 [2,3] 是 2 维,一张彩色图 [3,224,224] 是 3 维,一组图 [4,3,224,224] 是 4 维。concat 要求参与操作的张量秩相同,并且除了拼接轴以外的所有维度大小相同,结果秩不变。它做的是“扩容”,把某个已有维度的大小变大。stack 则要求所有输入张量的形状完全相同,结果秩加一,做的是“加维”。
举个例子:两个 shape 为 [2,3] 的矩阵,concat 之后可以是 [4,3] 或 [2,6],秩仍然是 2;stack 之后可以是 [2,2,3]、[2,2,3] 或 [2,3,2],秩变成 3。这两个操作在最终结果上很容易分辨:只要结果比输入多了一个维度,那就一定是 stack 或 stack 的等价形式;如果维度数量没变,只是某一维变大,那就是 concat。
2.2 好比“挂车厢”和“摞盘子”
我用两个生活化类比来记。concat 像给列车挂车厢:列车本身的维度不变,还是一列车,但长度变长了。stack 像把盘子一个个摞起来:每个盘子还是原来的形状,但新增了“第几层”这个概念。如果你把两个操作名换着用,就像把车厢并排焊在一起,或者把盘子一个接一个平铺,语义完全乱掉。
这个类比还能帮你理解 axis/dim 的位置。连接车厢时,你可以选择挂在车头后面还是车尾后面,这就是 concat 的拼接轴;摞盘子时,你可以决定“层编号”放在最前面、中间还是最后,这就是 stack 的新维度插入位置。
2.3 对比表格和关键结论
| 对比项 | concat / concatenate | stack |
|---|---|---|
| 核心行为 | 沿已有维度合并 | 插入新维度后堆叠 |
| 结果维度(秩) | 不变 | 加 1 |
| 输入形状要求 | 除拼接维度外完全相同 | 所有张量形状完全相同 |
| 典型用途 | 特征拼接、合并 batch | 从样本列表构造 batch |
| 新维度位置 | 无,只能是已有维度 | 由 axis/dim 指定 |
关键结论一句话:判断用哪个,先看最终张量是不是比每个输入都多一个维度。多,就用 stack;不多,就用 concat。这条规则几乎能解决所有选型问题。
3. 三大库 API 对比:PyTorch、TensorFlow、NumPy 的写法差异
3.1 同一个例子,三种写法
为了让对比直观,我统一用两个 shape 为 [2,3] 的随机张量做实验。下面分别是 NumPy、PyTorch、TensorFlow 的写法,结果形状完全一致。
# NumPy import numpy as np a = np.random.randn(2, 3) b = np.random.randn(2, 3) c_cat0 = np.concatenate([a, b], axis=0) # [4, 3] c_cat1 = np.concatenate([a, b], axis=1) # [2, 6] c_stack0 = np.stack([a, b], axis=0) # [2, 2, 3] c_stack1 = np.stack([a, b], axis=1) # [2, 2, 3] c_stack2 = np.stack([a, b], axis=2) # [2, 3, 2]# PyTorch import torch a = torch.randn(2, 3) b = torch.randn(2, 3) c_cat0 = torch.cat([a, b], dim=0) # [4, 3] c_cat1 = torch.cat([a, b], dim=1) # [2, 6] c_stack0 = torch.stack([a, b], dim=0) # [2, 2, 3] c_stack1 = torch.stack([a, b], dim=1) # [2, 2, 3] c_stack2 = torch.stack([a, b], dim=2) # [2, 3, 2]# TensorFlow import tensorflow as tf a = tf.random.normal([2, 3]) b = tf.random.normal([2, 3]) c_cat0 = tf.concat([a, b], axis=0) # [4, 3] c_cat1 = tf.concat([a, b], axis=1) # [2, 6] c_stack0 = tf.stack([a, b], axis=0) # [2, 2, 3] c_stack1 = tf.stack([a, b], axis=1) # [2, 2, 3] c_stack2 = tf.stack([a, b], axis=2) # [2, 3, 2]注意 a、b 是 [2,3] 时,stack(dim=0) 和 stack(dim=1) 的结果形状都是 [2,2,3],但内容不一样。比如 a 的第一行是 [1,2,3],b 的第一行是 [7,8,9]。stack(dim=0) 会得到两个“平面”,第一个平面是 a,第二个平面是 b;stack(dim=1) 得到的第一个平面由 a、b 的第一行组成,第二个平面由 a、b 的第二行组成。所以不要只看 shape,还要看你想把“第几个张量”作为新维度的索引。
3.2 API 细节差异
这些细节最容易在写代码时被忽略:
- np.concatenate 的 axis 可以省略,默认是 0;np.stack 的 axis 默认也是 0。
- torch.cat 的 dim 是必填的,不写会直接报错;torch.stack 的 dim 默认是 0,但我建议每次都显式写出来。
- tf.concat 的 axis 必填,tf.stack 的 axis 也建议显式给出。
- NumPy 还有 np.vstack、np.hstack、np.dstack 这些快捷函数。它们本质上是 concatenate 的封装,并不是 stack。vstack 会把一维数组先变成二维再沿第一维拼,命名确实有迷惑性。
三个库都支持负索引。axis=-1 表示最后一个维度。比如 torch.stack([a,b], dim=-1) 在 a、b 为 [2,3] 时等价于 dim=2,得到 [2,3,2];np.concatenate([a,b], axis=-1) 等价于 axis=1,得到 [2,6]。在处理不知道尾部形状时,负索引很方便,但先确认自己理解它在当前 shape 下指向哪里。
3.3 TensorFlow 动态维度要注意
TensorFlow 的 shape 里经常有 None,表示动态 batch 大小。tf.concat 和 tf.stack 都要求非拼接维度的静态 shape 能推断出来,否则会在构建图时报错。比如两个 shape 为 [None, 3] 的张量做 concat(axis=0),结果是 [None, 3];做 concat(axis=1),结果是 [None, 6];做 stack(axis=0),结果是 [2, None, 3](新维大小固定为 2)。如果你用了 Dataset 的 padded_batch,生成的张量可能带 None,这时更要留意 axis 的选择。
4. 为什么 Stack 等价于“先升维再 Concat”:拆开看原理
4.1 关键等价式
stack 并不是什么黑魔法,它在概念上可以拆成两步:先对每个输入张量做 expand_dims / unsqueeze,在目标位置插入一个大小为 1 的维度,然后再沿这个新维度做 concat。用 NumPy 表示:
np.stack([a, b], axis=k) == np.concatenate( [np.expand_dims(a, k), np.expand_dims(b, k)], axis=k )用 PyTorch 表示:
torch.stack([a, b], dim=k) == torch.cat( [a.unsqueeze(k), b.unsqueeze(k)], dim=k )这个等价不是实现细节,而是对“新增维度”的严格定义。如果我们要把两个完全一样的张量变成一个更高维的张量,唯一合理的办法就是让它们先各带一个大小为 1 的新轴,然后在这个新轴上合并。stack 只是把这两步封装成了一个函数。
4.2 用具体 shape 推一遍
以 a、b 都是 [2,3] 为例:
- k=0:unsqueeze(0) 变成 [1,2,3],cat axis=0 得到 [2,2,3],新维度在最前面。
- k=1:unsqueeze(1) 变成 [2,1,3],cat axis=1 得到 [2,2,3],新维度在中间。
- k=2:unsqueeze(2) 变成 [2,3,1],cat axis=2 得到 [2,3,2],新维度在最后。
所以 stack(dim=0) 和 stack(dim=1) 虽然最终 shape 一样,但内部数据排布完全不同。看下面的小例子:
import torch a = torch.tensor([[1, 2, 3], [4, 5, 6]]) b = torch.tensor([[7, 8, 9], [10, 11, 12]]) s0 = torch.stack([a, b], dim=0) # shape [2,2,3],第一个“平面”是 a,第二个“平面”是 b # [[[1,2,3],[4,5,6]], [[7,8,9],[10,11,12]]] s1 = torch.stack([a, b], dim=1) # shape [2,2,3],第一个“平面”由第一行组成 # [[[1,2,3],[7,8,9]], [[4,5,6],[10,11,12]]]如果后续要按“第几个原始张量”索引,用 stack(dim=0);按“同一行的不同张量”索引,用 stack(dim=1)。这在多视图、多分支特征合并时经常出现。
4.3 expand_dims 是视图,stack 会拷贝
更准确地说,NumPy 和 PyTorch 里的 expand_dims / unsqueeze 通常返回共享底层数据的视图,不会立刻复制内存。真正的复制发生在 concatenate / cat 阶段,因为目标张量需要连续存放不同输入的数据。理解了这一点,就能明白 stack 不是简单地在原张量上增加一个轴,它必须产生一个全新的大张量。
这对性能优化很重要:不要把大量小张量在循环里反复 stack。每次 stack 都会发生一次性全量拷贝,如果循环累积,复杂度会接近 O(n²)。正确做法是先收集到一个 list,最后统一 stack 或 cat。我后面会再讲这个坑。
5. 实际场景选型:特征拼接、Batch 构建、序列处理分别该用谁
5.1 特征拼接用 concat
在深度学习模型里,最常见的“沿特征维合并”是多模态融合和特征金字塔。假设两个 feature map 形状都是 [N,256,16,16],你想把两个分支的特征合并成 [N,512,16,16],用的是 torch.cat(features, dim=1),PyTorch 中通道维是 1。如果误用 stack,会得到 [2,N,256,16,16] 或 [N,2,256,16,16],语义完全变了:你不再增加通道数,而是把“分支个数”作为一个新维度保留。某些注意力模型里确实需要这种 group 维度,但普通卷积层通常不能直接处理。
f1 = torch.randn(8, 256, 16, 16) f2 = torch.randn(8, 256, 16, 16) fused = torch.cat([f1, f2], dim=1) # [8, 512, 16, 16]同理,两段文本的 embedding 是 [B,L,D] 和 [B,L,D2],沿特征维 concat 得到 [B,L,D+D2];沿时间维 concat 则要求其他维度相同。选哪个不是看个人喜好,而是看目标维度的语义是什么。
5.2 样本列表变 batch 用 stack
这是最核心的场景。一个 list 里的每张图是 [3,224,224],要变成 [N,3,224,224] 就必须用 stack。stack 也会出现在更复杂的结构里,比如对比学习中一个样本有多个数据增强后的视图,每个视图是 [3,224,224],你想保留“样本”和“视图”两个语义层次,就在 dim=1 上 stack:
views = [] for aug in augmentations: v = aug(image) # 每个 [3,224,224] views.append(v) batch_views = torch.stack(views, dim=1) # [N, num_views, 3, 224, 224]注意这里输入 views 的顺序和数量,会直接决定新维度的内容。stack 之后,索引 batch_views[0, 0] 是第一张图的第一个视图,batch_views[0, 1] 是第一张图的第二个视图。
5.3 合并已有 batch 用 concat
两个 batch,每个是 [8,3,224,224],合并成 [16,3,224,224] 用 cat(dim=0)。如果你想把两个 batch 组成一个“batch 的 batch”,才会用 stack 得到 [2,8,3,224,224]。分布式训练或者多卡推理时,需要把多个设备的结果聚合,通常也是 concat。先想清楚你要的结果是“样本数相加”还是“多了一个设备维度”,然后再选。
5.4 变长序列怎么办
stack 严格要求形状完全相同,所以变长序列不能直接 stack。例如三个句子的 token 数不同,直接 stack 会报错。标准做法是先 pad 到相同长度,然后用 stack 堆成 batch,另外维护一个 mask 表示真实长度。TensorFlow 的 Dataset.padded_batch 也在内部做类似的事:pad 之后实际上就是 stack 或 concat。所以当你看到 padded_batch 报维度错误时,也要往这个方向排查。
| 目标结果 | 操作 |
|---|---|
| [..., D1+D2] 不新增维 | concat |
| [2, ...] 最前面多一维 | stack(dim=0) |
| [..., 2, ...] 中间多一维 | stack(dim=k) |
| [..., 2] 最后多一维 | stack(dim=-1) |
6. 容易踩的坑和排查思路:维度报错、dtype 不一致、空张量
6.1 维度报错:看异常消息里的关键信息
PyTorch 的 torch.cat 报错通常是“Sizes of tensors must match except in dimension N...”。意思就是除了第 N 维(拼接的那个维度),其他维度的形状必须一致。比如你对 shape [2,3] 和 [2,4] 做 cat(dim=0),前者的列数是 3,后者是 4,就会报错。TensorFlow 常见的是“ConcatOp : Dimension X in both shapes must be equal”,NumPy 是“all the input array dimensions except for the concatenation axis must match exactly”。这三句话几乎是同一个意思,看到它们就说明你正在 concat,而不是 stack。
如果 stack 报错,通常是“all input tensors must have the same shape”,因为 stack 要求所有输入完全相同。这时候要去检查数据读取流程,看是不是某张图缩放过、通道数不一样,或者前一步不小心做了个多余操作。
6.2 dtype 不一致:有的框架会悄悄帮你转型
NumPy 的 np.concatenate 如果传 int32 和 float64,会自动提升为 float64,一般不报错,但结果可能不是你要的精度。PyTorch 的 torch.cat 要求所有输入 dtype 一致,不一致会直接报 TypeError。TensorFlow 2 默认也要求 dtype 一致。从 NumPy 习惯迁移过来的人,很容易在 PyTorch 里突然被 dtype 报错卡住。排查方式很简单:打印每个张量的 dtype、shape、device。比如:
def debug_tensors(tensors, name="tensors"): for i, t in enumerate(tensors): print(name, i, t.shape, t.dtype, t.device)device 不一致也会报错,需要先统一 .to(device)。在 GPU 训练时,如果某个样本来自 CPU 上的 NumPy 数组忘记转成 CUDA 张量,cat 和 stack 都会报错。
6.3 空张量和列表为空
torch.cat([]) 会抛出 RuntimeError,提示 expected a non-empty list。np.concatenate([]) 也一样。在写动态汇总逻辑时,如果待拼接的列表可能为空,不要直接 concat 或 stack。可以先判断:
if len(tensors) == 0: return torch.zeros(0, feature_dim, device=device)另外,如果张量里混入零大小维度,比如 [0,3,224,224],cat 可能得到空结果但不出错,stack 则可能因为新维大小固定而表现不同。这类边界行为在不同框架里并不完全一致,调试时要格外小心,最好打印出来确认。
6.4 循环里反复 cat 的性能陷阱
不要在 for 循环里不断 torch.cat([acc, x], dim=0),因为每次 cat 都会分配一个新张量并把之前所有内容拷贝一遍,复杂度接近 O(n²)。对视频帧、流式数据、动态图尤其明显。正确做法是先用 list 收集所有结果,循环结束后一次性 cat 或 stack;如果预先知道总数,可以预分配 [N,...] 的张量,然后按索引赋值。下面是一个对比:
# 不推荐:循环里反复 cat acc = torch.zeros(0, 10) for x in frames: acc = torch.cat([acc, x.unsqueeze(0)], dim=0) # 推荐:先收集再一次性拼接 buf = [] for x in frames: buf.append(x.unsqueeze(0)) result = torch.cat(buf, dim=0)这个习惯同样适用于 NumPy 和 TensorFlow。我在处理大量小片段时经常看到这个坑,改完后速度提升非常明显。
6.5 梯度与计算图
torch.cat 和 torch.stack 都是可微操作,会正常建立 autograd 图。如果只是做数据预处理,可以放在 torch.no_grad() 里,避免给不需要梯度的张量留下额外图信息。在模型前向里用 stack 把多个 head 的输出合并时,梯度会流向所有参与计算的输入。如果某个输入不需要梯度,记得先 detach。TensorFlow 的 GradientTape 同理。大多数场景不需要过度担心,但当你发现显存异常上涨时,可以检查是不是在计算图里堆了太多中间张量。
我个人现在的习惯是:写任何拼接之前,先在纸上写出“输入 shape -> 输出 shape”,明确目标是“多一维”还是“某一维变长”。这个简单的检查,比任何 API 记忆都可靠。PyTorch、TensorFlow、NumPy 之间切换时,我几乎都是靠这个判断来避免出错。如果你也经常被 Stack 和 Concat 搞晕,可以试试这个方法,应该会少写很多 debug 代码。