最近在整理深度学习张量操作的系列笔记时发现一个特别有意思的现象很多基础不错的开发者一说起Stack和Concat的区别嘴上能说出一个加维度一个不加维度但真到了写代码的时候尤其是要从PyTorch切到TensorFlow或者NumPy或者在多头注意力机制里处理head维度时照样会翻车。甚至有人会把torch.stack当成torch.cat用结果训练到一半遇到shape不匹配的报错一查原来是Attention Mask的维度多了一层或者少了一层——这种问题一旦出现排查起来格外费劲。这篇文章我就专门把这三个框架里的Stack和Concat掰开揉碎讲清楚。核心内容来自我做的PyTorch - 28 - PyTorchTensorFlow和NumPy中的Stack vs Concat - 深度学习Tensor Ops这期笔记我把它整理成一篇完整的手册式文章覆盖概念拆解、三个框架的API对比、实际场景下的选择逻辑以及我在坑里爬出来后整理的排查技巧。不管你是刚入门Transformer的新手还是已经在写分布式训练脚本的老手这篇文章都值得花十五分钟看看。1. 先搞清楚Stack和Concat在维度上的本质区别很多教程喜欢用Axis就是第几维来解释这个问题但我发现单纯讲axis参数的话新手记住后面马上就忘。我习惯用一个更生活的类比Concat是把多段绳子接起来绳子本身的粗细维度数不变只是变长了Stack是给每根绳子套上一个新的包装盒然后把这些独立的盒子摞在一起盒子的数量变成了新的一个维度。1.1 Concat的维度逻辑详解Concat在PyTorch中叫torch.cat在TensorFlow中叫tf.concat在NumPy中叫np.concatenate做的事情是沿着某个已有维度把所有元素拼接起来。关键点在于——拼接前后的总维数相同。用实际形状来举例会更清晰一些。假设现在有一批序列特征形状是[batch_size, seq_len, hidden_size]比如[32, 128, 768]。如果我把这一批数据沿着seq_len这一维拼接也就是axis1把两条数据拼在一起最终得到的形状是[32, 256, 768]——第三个维度参数维度不变但序列长度翻倍了。这正是Concat最常见的用途在处理变长序列批量计算时尤其重要。再举一个更具体的例子。在自注意力机制里要把query和key拼接起来做后续计算如果query的形状是[batch, heads, seq_q, dim]key的形状是[batch, heads, seq_k, dim]那么torch.cat([q, k], dim2)得到的就是[batch, heads, seq_q seq_k, dim]。这个时候heads维度完全没有被动过这个非常关键。1.2 Stack的维度逻辑详解Stack做的事情核心是把多个张量当成独立个体增加一个新的维度来盛放这些个体。比如我有三个形状都是[seq_len, hidden]的张量torch.stack([tensor1, tensor2, tensor3], dim0)结果形状变成[3, seq_len, hidden]。注意看这里凭空多了一个维度3出来而原始的数据维度被推到了后面的位置上。《A ConvNet for the Year》那篇论文里有的结构我现在记不太清了但有一件事情我记得特别牢——在处理多头注意力时把单头的结果堆叠成多头张量用的就是torch.stack。每头输出的[batch, seq, head_dim]通过torch.stack(heads, dim2)后合并成[batch, seq, num_heads, head_dim]。这个num_heads就是新加出来的维度。所以我说Stack是用来创造维度的Concat是用来扩充已有维度的。用表格来做一次直观对比操作维度数量变化例子输入两个[3, 4]张量输出形状记忆口诀Concataxis0不变两个[3,4]沿第一维拼接[6, 4]接绳子Concataxis1不变两个[3,4]沿第二维拼接[3, 8]加宽Stackaxis0增加一维两个[3,4]堆成列表[2, 3, 4]装箱Stackaxis1增加一维两个[3,4]插入新维[3, 2, 4]装箱插位这个表格值得收藏基本能解决90%的日常困惑。1.3 一个容易混淆的规则axis的取值含义许多人第一次接触Stack时最迷惑的点是dim或axis到底指向哪里。我总结一个非常实用的观察方法看输出张量中被新插入的维度在哪个位置。拿两个形状都是[3, 4]的张量举例np.stack([a, b], axis0)结果[2, 3, 4]。意思是新增维度在最前面堆的行为发生在第0号位置之前。np.stack([a, b], axis1)结果[3, 2, 4]。新增维度插到了中间每个小元素的第1个位置前塞了一个索引面板。np.stack([a, b], axis2)结果[3, 4, 2]。新增维度落在了最后相当于把两个矩阵的每个元素配对收进一个二元组里。这个规则在PyTorch中也完全一致torch.stack里的dim参数同样表示新维度要插入的坐标位。这个规则如果悟透了三维以上张量的操作就再也不会犯迷糊了。2. 三个框架下的Concat实现细节对比既然模块化开发免不了在几个框架之间来回切换那么熟悉每一个框架的API细节尤其是那些容易忽略的参数和陷阱就很值得花时间。2.1 NumPy中的np.concatenateNumPy是整个Python科学计算生态的地基理解了np.concatenate其他框架的concat基本就是换个名字的事。np.concatenate的基础用法import numpy as np a np.array([[1, 2], [3, 4]]) b np.array([[5, 6], [7, 8]]) # 沿着第0维行方向拼接 c np.concatenate([a, b], axis0) # 输出形状为 (4, 2) # 沿着第1维列方向拼接 d np.concatenate([a, b], axis1) # 输出形状为 (2, 4)这是我见过的最基础示例。但我想提醒一个极易出错的点np.concatenate要求所有输入张量在非拼接维度上的形状必须完全一致。比如上面这个例子如果a的形状是[3, 2]而b的形状是[4, 2]它们仍然可以沿第0维拼接因为列数都是2。但如果你试图沿axis1拼a的行数是3而b的行数是4分母不一致NumPy会直接抛出ValueError。这个机制在深度学习里表现为拼接embeddings和positional_encoding时seq_len维度不一致必报错。2.2 PyTorch中的torch.cattorch.cat从接口上来看跟np.concatenate几乎一致但有几个细节需要额外留意。首先是参数位置的差异。torch.cat的第一参数是一个张量列表tuple或list第二个参数是dim默认值是0。torch.cat在dim取值上比NumPy多了些灵活性支持负数索引比如dim-1表示最后一个维度。import torch # 两个形状为 [2, 3] 的张量 x1 torch.randn(2, 3) x2 torch.randn(2, 3) # 沿第0维拼接得到 [4, 3] y0 torch.cat([x1, x2], dim0) # 沿最后一维拼接得到 [2, 6] y_last torch.cat([x1, x2], dim-1)dim-1这种写法在注意力机制里非常常用。比如合并多个特征向量时我经常用torch.cat([feature1, feature2], dim-1)把两个向量拼成一个长向量再接全连接层。torch.cat和torch.stack在PyTorch里是两个不同的函数这一点要比TensorFlow更明确因为TensorFlow中也有tf.concat和tf.stack两者的语义是完全一致的但TensorFlow对动态shape的支持会好很多。PyTorch里还有一个容易踩的点torch.cat要求所有张量必须在同一设备上CPU或GPU否则会报错RuntimeError: Expected all tensors to be on the same device。我在做多卡训练时遇到过一次一个tensor在cuda:0另一个在cuda:1一cat就报错。所以写多卡代码时concat前一定要确认设备位置。2.3 TensorFlow中的tf.concattf.concat的接口形如tf.concat(values, axis)。跟PyTorch的区别在于——它的axis是必须显式指定的没有默认值。如果你忘了传axis会直接报错TypeError。TensorFlow有个特点比较友好支持axis为负数例如tf.concat([a, b], axis-1)同时支持axis是Python列表例如tf.concat([a, b, c], axis[-1])不这个不能随便乱用因为tf.concat的axis在低版本中支持整数新版本里虽然可以传入张量但语义上仍然是整数方向。千万不要被这种写法迷惑。真正值得说的是TensorFlow 2.x 默认是动态图模式Eager Execution所以在tf.concat中拼接tensor时张量的形状是实时可知的。但如果你用tf.function装饰了函数就要特别小心tf.concat中的axis不能被张量化必须是Python整数或者标量张量。我曾经在tf.function里用了一个来自tf.shape的动态值作axis结果直接报错Operations that dont support gradients are not allowed——虽然错误提示不一定是这个但排查了很久才发现是axis必须静态化。回顾一下在三种框架里Concat的语义是完全统一的但API的细节差异设备限制、动态值限制、默认参数都会影响实际使用。所以在切换框架时不要只盯着函数名还要注意框架的运行时策略。3. 三个框架下的Stack实现细节对比Stack在三个框架里的API形态和维度逻辑也是一脉相承的但各自有不同的坑位。3.1 NumPy中的np.stacknp.stack接受一个数组序列在指定轴上堆叠。最典型的场景是把一组向量变成矩阵比如把多个[embedding_dim]的向量堆成[num_sentences, embedding_dim]。import numpy as np emb1 np.array([0.1, 0.2, 0.3]) emb2 np.array([0.4, 0.5, 0.6]) emb3 np.array([0.7, 0.8, 0.9]) # 堆叠成 [3, 3] batch np.stack([emb1, emb2, emb3], axis0)这个在训练循环里构建batch的时候特别好用循环里收集了一堆形状相同的样本最后用np.stack一次性堆成batch。相比之下如果用np.concatenate你得先把每个样本unsqueeze扩维一次再拼接绕了一圈还是回到了Stack的语义上。np.stack并非没有坑。它要求列表中所有数组的shape完全一致。这里说的完全一致比np.concatenate更严格因为stack会新增一维它必须保证新增维度之外的所有维度都对齐。如果有一个数组是[3]另一个是[4]np.stack会直接报错没有任何商量的余地。3.2 PyTorch中的torch.stacktorch.stack和torch.cat在PyTorch中共享一个存储布局逻辑但torch.stack会创建新的副本也就是说它将每个输入张量复制之后再放到新的维度上。这一点从内存占用上就很明显如果输入N个张量输出占用的内存大约是输入总和加上新维度的元数据开销。import torch a torch.tensor([1, 2, 3]) b torch.tensor([4, 5, 6]) # 变成 [2, 3] ab_stack torch.stack([a, b], dim0) print(ab_stack.shape) # torch.Size([2, 3]) # 变成 [3, 2] ab_stack_1 torch.stack([a, b], dim1) print(ab_stack_1.shape) # torch.Size([3, 2])如果你注意观察dim0时a变成了结果的第一行b变成了第二行dim1时a和b变成了结果的每一行各自的前两个元素。这正是新维度插到哪里的直观体验。PyTorch中关于torch.stack的一个常见陷阱出现在梯度计算中。某些老版本下torch.stack对输入的梯度传播可能出现问题尤其是当你多次stack并且传入requires_gradTrue的张量时。虽然新版本已经修复但为了保险起见如果反向传播中出现奇怪的梯度NaN可以试试用torch.stack([a, b], dim...)替换成torch.cat加unsqueeze的组合。3.3 TensorFlow中的tf.stacktf.stack在TensorFlow中的行为跟PyTorch基本一致但它有一个独特的优势对动态形状动态shape的友好处理。import tensorflow as tf a tf.constant([[1, 2], [3, 4]]) b tf.constant([[5, 6], [7, 8]]) # 堆叠为新维度在第一维 c tf.stack([a, b], axis0) # 形状 [2, 2, 2] # 堆叠为新维度在第二维 d tf.stack([a, b], axis1) # 形状 [2, 2, 2]咦axis0和axis1输出形状看起来一样那二者的区别在哪里这两个输出的形状确实一样都是[2,2,2]但数据排布不同。c里面第0维是第几个输入张量d里面第1维是第几个输入张量。如果要具体解释可以想象成c是两页纸叠在一起每页上是2行2列的表格而d是每个表格行内多了一个来源索引列的区别——不对还是直接用索引取值来理解更准确c[0]是ac[1]是b而d[:, 0, :]是a的第一行和b的第一行交替组成这样子确实容易绕。我建议遇到这种困惑时就老老实实写一行测试代码 c[0] tf.Tensor: shape(2, 2), dtypeint32, numpy array([[1, 2], [3, 4]])一句话总结先搞清楚哪个轴代表第几个输入张量再谈其他细节。tf.stack在tf.function中同样存在轴必须静态化的要求跟tf.concat一样axis不能是动态张量必须是Python常量。4. 实际场景中的选择策略与优先级了解了工具本身还得知道什么业务场景该用哪个。我根据自己的实战经验把几个高频场景的选型逻辑整理一下。4.1 场景一多头注意力机制中的head合并这是面试里高频考点也是真正在Transformer代码里出现频率最高的场景。在标准多头注意力中通常会先经过一个线性层把hidden_size分成num_heads * head_dim然后reshape到[batch, num_heads, seq_len, head_dim]。接下来每个头独立计算注意力。在合并各个头的结果时有两种实现路径路径Atorch.cat配合reshape。 先把所有头的输出沿着head_dim拼接变成[batch, seq_len, num_heads * head_dim]再通过一个线性层做混合。这种写法用的是Concat的逻辑因为头的数量不需要额外维度来表示本来就藏在特征维度里。路径Btorch.stack后平均。每个头输出[batch, seq_len, head_dim]用torch.stack(head_outputs, dim2)汇总成[batch, seq_len, num_heads, head_dim]再mean或者weighted sum。这种写法多了一个stack维度后续若要mean就直接在num_heads维度上约简。我自己的习惯是如果是标准Transformer的实现几乎不会用到Stack因为头已经reshape进特征维度了如果你在处理类似多个独立模型的集成这类工作Stack才真正发挥威力。4.2 场景二构建训练Batch这个场景对很多初学者很友好也是区分Stack和Concat的绝佳例子。假设我在循环里为每个样本单独计算了特征向量每个特征向量的形状是[feature_dim]。我要把它变成一个batch供模型处理features [] for i in range(batch_size): feat model(sample_i) # 形状 [feature_dim] features.append(feat) # 错误做法np.concatenate(features, axis0) - 得到 [batch * feature_dim]全糊在一起了 # 正确做法np.stack(features, axis0) - [batch, feature_dim] batch_feature np.stack(features, axis0)为什么不能用concat因为单个样本的特征向量只有一维concat只能把它们首尾相接变成一个大长条这个长条不再是不同样本、同一特征位的二维结构。而stack天然能把多个一维向量变成二维矩阵每一行是一个样本每一列是一个特征位。在日常编码里如果我先收集了多个[seq, dim]的序列想把它们变成一个batch也需要用stack。这个习惯一定要养成。4.3 场景三拼接embedding与额外特征推荐系统或者多模态模型中经常要把文本的embedding和手工特征、图像特征拼起来。这时用的是Concat原因是多个来源的特征共享同一个batch与seq_len维度只是特征维度不同。比如text_emb model(text_tokens) # [batch, num_tokens, 768] img_feat cnn_model(images) # [batch, 2048] img_feat_expanded img_feat.unsqueeze(1).expand(-1, num_tokens, -1) # 强行扩到 [batch, num_tokens, 2048] combined torch.cat([text_emb, img_feat_expanded], dim-1) # [batch, num_tokens, 7682048]这个expand步骤很容易忘。如果不展成相同的seq_lentorch.cat会直接报错因为非拼接维度的长度必须相等。我在多模态项目里第一次写这个逻辑时就被卡住了后来才反应过来需要先把那个batch, 2048的特征expand到每个token位置上。所以在动态句长和静态图像特征共存的场景下Concat之前往往需要做维度广播或repeat这步很容易漏。4.4 场景四动态计算图中的形状推断在TensorFlow的动态图Eager模式里tf.stack会在运行时动态计算输出形状这在构建复杂输入流水线时极为方便。但在tf.function静态图模式下tf.stack的输入个数和形状必须是可推断的否则会报错。我建议如果确实需要在动态张量数量下例如循环中不确定会产生多少特征使用tf.stack最好先使用tf.TensorArray收集张量最后调用.stack()方法。TensorArray.stack()内部还是会用栈的方式把列表合并成张量语义一样但允许在循环中动态写入。ta tf.TensorArray(dtypetf.float32, size0, dynamic_sizeTrue) for i in range(num_steps): feat model(...) # [batch, dim] ta ta.write(i, feat) all_feat ta.stack() # [num_steps, batch, dim]这种写法在循环内部非常实用既避免了Python列表的低效又不会有静态图推断的问题。5. 踩坑实录与排查技巧这部分算是我这么多年实实在在趟出来的经验。每次写代码遇到Stack和Concat相关的报错基本都在下面这几类里。5.1shape不匹配的报错这类报错最常见也最容易被忽略。以torch.cat为例RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 5 but got size 6 for tensor number 1 in the list.这个报错非常直白除了dim1拼接的那个维度可以不同其他维度必须严格一致。遇到这种情况我第一反应是打印出每个输入张量的shape直接看一遍是哪一个维度对不上。实战中遇到的一次是BERT的attention_mask形状是[batch, seq_len]而token_type_ids形状是[batch, seq_len - 2]因为我在预处理时对token_type_ids做了截断忘了补回来。排查了十分钟最后把token_type_ids做了一次pad才解决。5.2 设备不一致导致的神秘报错PyTorch中torch.cat对设备一致性要求极严。你可能在CPU上做数据预处理时生成了一个mask但模型输入在GPU上直接cat就会报错。报错文案为RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!解决方案很简单显式调用.to(device)比如mask mask.to(device) combined torch.cat([input_ids, mask], dim-1)我还建议在数据加载器里把所有输入统一放到device上之后再进模型而不是在模型内部散落地到处to(device)这样出问题也容易排查。5.3 梯度传播中的Stack问题在这个问题上PyTorch相对稳定TensorFlow要注意更多。TensorFlow中如果你在一个tf.function内对tf.stack的结果做reshapereshape和stack的梯度可能会互相干扰。尤其是在使用混合精度训练mixed_precision时梯度可能因为堆叠操作而变成None或NaN。我遇到过一次损失函数里用tf.stack把多个对比损失堆起来然后加权求和结果低精度模式下梯度变成NaN。排查了很久最终把tf.stack改成tf.reduce_sum加expand_dims组合才稳定。5.4 在torch.stack和torch.cat之间来回切换的心法有时候看到别人代码里一会儿用cat一会儿用stack自己融不进去。我觉得最稳的判断方法就一句话**拼完之后shape变少没**如果拼完后维度数没变那就是Concat如果拼完后维度数多了一层那就是Stack。用这个标准再审视一遍想把[batch, seq, dim]和[batch, seq, dim]拼成[batch, seq, 2*dim]维度数还是3用cat(dim-1)。想把两个[batch, seq, dim]汇总为[2, batch, seq, dim]——维度数从3变成4用stack(dim0)。这个心法可以应对95%的场景。5.5 性能层面的注意事项关于性能我观察到的现象是stack通常比cat稍慢因为它涉及新维度的创建和数据的重新排布。在多卡大模型场景下尤其要警惕stack带来的额外显存拷贝。如果一个操作仅仅是为了把多个张量并排放到一个维度上用cat更划算如果必须增加维度也别无选择只能用stack。还有一个已知的优化技巧在PyTorch中如果先unsqueeze再cat与直接stack在功能上等价但stack在底层有专门的内核优化速度反而更快。因此不要为了避开stack而手动扩维再cat直接用torch.stack更省事。6. 一个综合示例三种框架实现同一任务为了帮助你把知识串起来我设计了一个小小的任务把3个形状为[2, 3]的矩阵在三种框架下分别实现沿着行的方向堆叠成[3, 2, 3]以及按行拼接成[2, 9]。这个任务做起来非常能检验理解程度。6.1 NumPy版本import numpy as np a np.array([[1, 2, 3], [4, 5, 6]]) b np.array([[7, 8, 9], [10, 11, 12]]) c np.array([[13, 14, 15], [16, 17, 18]]) # Stack - [3, 2, 3] stacked np.stack([a, b, c], axis0) print(stacked.shape) # (3, 2, 3) # Concat - [2, 9] concatenated np.concatenate([a, b, c], axis1) print(concatenated.shape) # (2, 9)注意axis的选择在np.concatenate里是沿着哪个维度进行拼接在np.stack里是新维度插到哪个位置。这两个语义虽然都用一个axis但在理解上不要搞混。6.2 PyTorch版本import torch a torch.tensor([[1, 2, 3], [4, 5, 6]]) b torch.tensor([[7, 8, 9], [10, 11, 12]]) c torch.tensor([[13, 14, 15], [16, 17, 18]]) stacked torch.stack([a, b, c], dim0) # [3, 2, 3] concatenated torch.cat([a, b, c], dim1) # [2, 9]PyTorch里dim0的stack和cat在形状上的差异是最容易被一眼看穿的。6.3 TensorFlow版本import tensorflow as tf a tf.constant([[1, 2, 3], [4, 5, 6]]) b tf.constant([[7, 8, 9], [10, 11, 12]]) c tf.constant([[13, 14, 15], [16, 17, 18]]) stacked tf.stack([a, b, c], axis0) # [3, 2, 3] concatenated tf.concat([a, b, c], axis1) # [2, 9]到这里你会发现只要能准确说出这个操作会让维度数量变吗任何框架都不会再成为障碍。7. 三个框架异同速查表把三种框架中Stack和Concat的API做一个直观对比方便遗忘的时候快速查功能NumPyPyTorchTensorFlowConcatnp.concatenate((a,b), axisk)torch.cat([a,b], dimk)tf.concat([a,b], axisk)Stacknp.stack([a,b], axisk)torch.stack([a,b], dimk)tf.stack([a,b], axisk)维度数量不变Concat不变cat不变concat维度数量增加Stack增加stack增加stack默认axisaxis0dim0必须显式给出负数索引支持支持支持动态图不支持支持Eager支持Eager静态图N/AN/Atf.function中axis需静态说句实在话有了这张表日常开发里再遇到这个cat该用哪个参数之类的问题直接查表即可没必要强行背下来。最后分享一个自己的小习惯每次写完涉及Stack或者Concat的操作我会立刻打印一下输出张量的shape然后用注释写上预期形状。这个习惯帮我拦下来无数次低级错误。而且在代码评审的时候如果你的PR里有清晰的shape注释同事读起来也会轻松很多。# 预期结果: [batch, seq_len, 2*dim] out torch.cat([a, b], dim-1)这行注释看上去简单在维护老模型的时候能省下大把时间。既然都看到这里了不妨打开编辑器随便建三个小张量分别试一下stack和cat对比它们的输出形状亲自体会一下维度变化的直觉。等你熟练了以后这些操作就会像呼吸一样自然再也不会搞混了。