Tensor.flatten(start_dim) 是我在 PyTorch 里用得最勤的接口之一尤其是搭 CNN、做多头注意力、处理各种多维特征图的时候几乎每一版模型代码里都会出现它。别看它只是一行.flatten()真正能把start_dim用明白、用对不踩形状坑的其实没那么多。这个接口解决的核心问题很简单把一个高维张量“摊平”成低维张量同时让你决定从哪个轴开始摊、摊到哪个轴为止。适合所有刚接触 PyTorch 的初学者、写模型时经常被维度搞晕的算法工程师以及想在 view、reshape、flatten 之间做出正确选择的同学。我见过太多人遇到形状对不上就直接x.view(x.size(0), -1)一把梭结果遇到非连续张量就报错也见过有人想保留 batch 维度做展平结果用x.flatten()把 batch 也揉进去了训练直接崩。这篇文章不会只给你 API 文档我会把维度怎么数、start_dim到底控制什么、flatten 底层和 view/reshape 的关系、以及真实场景里的用法全部拆开讲最后附上我实际踩过的坑。1. 先搞懂维度再学 start_dim1.1 维度不是玄学张量的轴到底怎么数想理解start_dim第一步不是背参数而是把“维度”这个词在脑子里变成一张立体的图。你可以把一个张量想象成一个多层书架第 0 维是书架的层数第 1 维是每一层上放的书第 2 维是每本书的页码。比如一个形状是(2, 3, 4)的张量就是 2 层书架、每层 3 本书、每本书 4 页。在 PyTorch 里dim0永远是最外层也就是shape元组里的第一个数dim1是第二层以此类推。很多人第一次接触时会把“列”和“行”搞混因为二维矩阵里dim0是行、dim1是列但一旦上升到三维、四维这种直觉就不够用了。我自己的习惯是每拿到一个张量先写一句注释把每个维度的含义标出来。比如一个图像 batch形状是(N, C, H, W)我就写# N: batch, C: channel, H: height, W: width。这一小步能帮你避开大部分维度错误。flatten这个操作本身不改变数据的总元素个数只改变张量的“排列层级”。它做的事情就是把指定的若干个连续维度合并成一个维度相当于把一个多层的嵌套结构压成一层。总元素个数等于各维度大小的乘积这个数字在展平前后是严格不变的。所以每次展平后你都可以用“总数守恒”的原则来验证自己的操作有没有错。1.2 start_dim 到底控制什么start_dim的全称是“起始维度”它决定从哪个轴开始合并。flatten 的规则是start_dim之前的维度全部保留从start_dim开始到end_dim结束的维度被合并成一个。默认情况下start_dim0也就是从最外层开始把所有维度合并成一个一维张量默认end_dim-1也就是一直合并到最后一维。举个例子形状为(2, 3, 4)的张量如果调用t.flatten(start_dim1)那第 0 维的 2 保留不动从第 1 维开始到最后一维合并结果是(2, 12)。这里 12 就是3 * 4。如果你调用t.flatten(start_dim0, end_dim1)那就是把前两个维度合并结果是(6, 4)。你看start_dim就像一个“钉子”它钉住了前面的维度不让动只让后面的维度合并起来。这个参数最常见的用途就是“保留 batch 维度”。在神经网络里batch是第 0 维它代表一次喂进去的样本数。当你要把特征图送入全连接层时不能把 batch 也展平进去否则全连接层会分不清哪些元素属于哪个样本。这时候写x.flatten(start_dim1)就保证了输出还是(batch_size, 特征长度)的形状。2. Tensor.flatten(start_dim) 的核心用法拆解2.1 一行代码全展平默认参数发生了什么先看最基本的调用方式。假设你有一个随机张量import torch t torch.randn(2, 3, 4) print(t.flatten().shape) # torch.Size([24]) print(t.flatten(0).shape) # torch.Size([24]) print(torch.flatten(t).shape) # torch.Size([24])三个调用的结果一样都是(24,)。因为start_dim默认是 0end_dim默认是 -1所以等价于把 2、3、4 三个维度全部合并成一维。这个操作在把输出特征变成一维向量时很有用比如你要对某个特征计算距离、做聚类或者把中间结果送入一个只接受一维输入的处理函数。这里要注意一个细节t.flatten()和t.flatten(0)虽然结果相同但语义上有一点点差别。flatten()用的是默认参数只要维度变化它都会重新计算范围flatten(0)是显式指定从第 0 维开始。如果你的张量维度个数是动态变化的显式写死某个数字并不总是安全。我一般会优先用默认参数或负数索引让维度的适应性更强。2.2 用 start_dim 保住 batch 维度实际写模型时用到最多的其实是start_dim1。考虑一个图像分类任务特征图经过卷积网络后形状是(8, 64, 7, 7)8 是 batch size64 是通道数7x7 是空间尺寸。你想送进全连接层就需要把通道和空间尺寸合并但 batch 必须留着。x torch.randn(8, 64, 7, 7) y x.flatten(start_dim1) print(y.shape) # torch.Size([8, 3136])3136 就是64 * 7 * 7。这一步做完之后y的形状正好是(batch, features)可以直接接nn.Linear(3136, num_classes)。如果你不小心用了x.flatten()得到的是(3136 * 8,)也就是(25088,)全连接层会直接报维度错误或者更隐蔽地如果你用x.flatten().view(batch_size, -1)这种操作虽然能恢复形状但多了一次不必要的复制纯属浪费。还有一个常见的等价写法是x.view(x.size(0), -1)。很多老代码都喜欢这么写它的效果和x.flatten(start_dim1)基本一致但view有一个前置条件张量必须是连续内存。关于这一点后面的第 3 节会详细讲。简单说从卷积层出来的特征图通常是连续的所以大多数时候view能用但一旦你把张量做了转置、切片、或者用transpose调整过维度view就会抛异常而flatten不会。2.3 end_dim 与负数索引更精细的控制flatten的第二个参数end_dim经常被忽略但它其实很有用。它控制展平的“终点”。比如你有一个形状为(batch, seq_len, hidden)的序列特征想只合并前两个维度保留最后一个维度x torch.randn(4, 16, 128) y x.flatten(start_dim0, end_dim1) print(y.shape) # torch.Size([64, 128])这在处理 RNN 输出时很常见RNN 的输出形状是(seq_len, batch, hidden)或(batch, seq_len, hidden)如果你想把它变成二维矩阵送入后续模块就需要用end_dim指定合并范围。更妙的是负数索引。start_dim-1表示从最后一维开始展平start_dim-2表示从倒数第二维开始展平。比如一个形状是(2, 3, 4, 5)的张量flatten(start_dim-2)会保留前两个维度合并最后两个维度结果是(2, 3, 20)。负数索引的价值在于当你的张量维度个数不确定时你依然可以表达“把最后两个维度合并”的意图不用写死绝对维度编号。我写通用工具函数时特别喜欢用负数索引因为代码在二维、三维、四维输入下都能正确工作。2.4 torch.flatten、Tensor.flatten 与 nn.Flatten 的区别PyTorch 里其实有三个带 flatten 字眼的东西很多人分不清torch.flatten(input, start_dim0, end_dim-1)是一个函数Tensor.flatten(start_dim0, end_dim-1)是一个张量方法两者等价只是调用姿势不同。还有一个nn.Flatten(start_dim1, end_dim-1)它是一个网络层模块。nn.Flatten有一个极其重要的默认值start_dim1。它故意把第 0 维也就是 batch 维保留下来因为绝大多数神经网络模块在使用展平时都不希望丢失 batch 信息。所以你在nn.Sequential里写nn.Flatten()时它等价于x.flatten(start_dim1)而不是x.flatten(start_dim0)。这个设计决策让很多新人在把torch.flatten用到网络里时踩坑你以为是同一个东西结果 batch 维被吞了。我建议在写网络层时统一用nn.Flatten()在写数据处理逻辑时用Tensor.flatten(start_dim...)这样语义清晰也更容易读代码。下面是三者对比调用方式默认 start_dim默认 end_dim典型使用场景torch.flatten(x)0-1把任意张量转成一维向量x.flatten(start_dim1)0-1保留 batch合并特征维度nn.Flatten()1-1在nn.Sequential中使用3. flatten、view、reshape 怎么选3.1 底层机制view、viewcontiguous、reshape 的真实差异写 PyTorch 的人早晚会遇到view、reshape、flatten这三个操作它们都能改变张量形状但底层的逻辑完全不同。view要求张量在内存中是连续的它只是改变了解释内存的方式不复制数据。如果张量不连续比如你对一个多维张量做了transpose再调用viewPyTorch 会直接抛出RuntimeError: view size is not compatible with input tensors size and stride。reshape是更宽松的版本如果张量连续它和view一样只改形状不复制如果不连续它会先复制一份连续内存再调整形状。这个“按需复制”的行为让reshape几乎不会报错但它可能带来一次额外的内存拷贝成本比view高一点。flatten的底层实现实际上是基于reshape的一个封装它内部先计算start_dim到end_dim之间所有维度的乘积然后调用reshape。也就是说flatten继承了reshape对非连续张量的宽容性。你可以放心地对转置后的张量做flatten它不会像view那样抛异常。3.2 内存连续性的坑内存连续性这个概念听起来抽象但它在实际开发里经常出来“咬你一口”。最常见的场景是transpose和permute之后x torch.randn(8, 3, 224, 224) x_t x.permute(0, 2, 3, 1) # 变成 (8, 224, 224, 3) # x_t.view(8, -1) # 这里会报错 y x_t.flatten(start_dim1) # 正常结果是 (8, 224*224*3)permute只是改变了遍历顺序的“视图”并没有真的在内存里把数据搬成新顺序。所以当你想把它展平时view就会因为内存布局不匹配而失败。flatten底层会先自动做contiguous再完成展平所以不会报错。代价是如果张量真的不连续它会触发一次数据复制。大多数情况下这个代价是值得的因为代码健壮性比那一点性能重要得多。一个小技巧如果你确定自己的张量是连续的又想省掉flatten内部可能的连续性检查可以先用x.contiguous().view(...)这样逻辑最明确。如果张量已经是连续的contiguous()是一个空操作几乎没有开销。3.3 工具选型建议结合我自己的经验给你一个简单粗暴的选择标准第一个只是改变形状不涉及维度的合并逻辑比如把一个(3, 4)变成(4, 3)转置用permute或transpose不要用 flatten。转置和展平是完全不同的语义只会搞混。第二个想把多维张量合并成一维或保留部分维度且不担心连续性问题用flatten(start_dim...)这是最优解。原因是它语义清晰、自动处理连续性、而且可读性比view(x.size(0), -1)好太多。view(x.size(0), -1)这种写法虽然老练但对不熟悉的人来说很容易误解为什么前面是x.size(0)后面是-1而flatten(start_dim1)一眼就能看出来是保留第 0 维、合并后续。第三个你明确知道张量连续并且追求极致的性能比如在高频循环里做形状变换可以用view省掉 reshape 内部的contiguous检查。但要写注释说明为什么这里安全。第四如果你在写一个对输入形状没有严格限制的通用模块优先flatten(start_dim1)它会兼容连续和非连续内存也几乎不会报错。4. 实际场景里的典型用法4.1 分类网络里接全连接层的展平CNN 分类网络的标准结构是“卷积特征提取 展平 全连接层”。卷积层的输出通常是四维(batch, channels, height, width)。全连接层nn.Linear接受的是二维输入(batch, features)中间必须有一次展平操作把后三维合并成一维。import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv nn.Sequential( nn.Conv2d(3, 16, 3, padding1), nn.ReLU(), nn.MaxPool2d(2) ) self.flatten nn.Flatten(start_dim1) self.fc nn.Linear(16 * 16 * 16, 10) def forward(self, x): x self.conv(x) # (batch, 16, 16, 16) x self.flatten(x) # (batch, 16*16*16) return self.fc(x)这里nn.Flatten(start_dim1)和内置的nn.Flatten()效果一模一样写出来是为了强调 start_dim 在这个场景里的作用。很多人会问为什么不能用全局平均池化代替展平可以但展平保住了所有空间信息池化会损失一部分。在需要保留细节的任务里展平是不可替代的。另外你完全可以在forward里手写x.flatten(start_dim1)但把它封装成层的好处是打印模型结构时更清楚。4.2 多头注意力中的维度重组多头注意力是 Transformer 里的核心模块它也是 flatten 的高频使用场景。Q、K、V 在拆分多头之后常见形状是(batch, heads, seq_len, head_dim)。有时候要把 batch 和 heads 合并起来一起参与矩阵乘法做完再拆回去。x torch.randn(4, 8, 16, 32) # (batch, heads, seq_len, head_dim) # 合并 batch 和 heads y x.flatten(start_dim0, end_dim1) print(y.shape) # (32, 16, 32)相当于把 4*8 个序列一起处理 # 恢复原形状 z y.view(4, 8, 16, 32) print(z.shape) # (4, 8, 16, 32)这种flatten加view的组合拳在实现 Grouped-Query Attention、Flash Attention 的各种变体里反复出现。还有另一种情况如果你想把(batch, heads, seq_len, head_dim)中的seq_len和head_dim合并让每个头变成二维矩阵可以用x.flatten(start_dim2)结果是(batch, heads, seq_len * head_dim)。这一步经常用于把注意力分数和特征矩阵做点积之前的预处理。4.3 特征融合与 embedding 输出在做多模态或特征融合时经常要把不同来源的特征拼在一起。假设一个分支输出三维特征(batch, seq_len, dim)另一个分支输出二维特征(batch, embedding_dim)它们维度不同无法直接 concat。通常做法是先让三维分支做 flatten 或池化变成二维再在特征维度上拼接。text_feat torch.randn(8, 16, 128) # (batch, seq_len, dim) image_feat torch.randn(8, 512) # (batch, img_dim) text_feat text_feat.flatten(start_dim1) # (8, 16*128) (8, 2048) fused torch.cat([text_feat, image_feat], dim1) # (8, 2560)这里如果你忘了start_dim1直接把text_feat展平成一维就会得到(8*16*128,)batch 维度消失拼接必然失败。另外一个非常常见的“送 embedding”场景是模型输出四维特征图(batch, channels, height, width)你想得到每个样本的 embedding 向量可以先做flatten(start_dim1)得到(batch, channels*height*width)再经过一层线性层投影到固定维度。这个操作在对比学习、人脸识别、图像检索里都是标配。4.4 根据概率张量采样一个值时flatten 的妙用前面提到的热词里有一个“torch 根据某个 tensor 来 sample 一个值”很多人做采样时也会被维度绊倒。最典型的是torch.multinomial它接收一个二维概率矩阵对每一行做一次多项分布采样返回索引。但如果你的概率分布是三维或四维的比如形状是(batch, height, width)的逐像素概率图想根据它采样一个像素位置直接采样会很麻烦。这时候 flatten 就派上用场了。先把概率图展平成(batch, height * width)然后对每一行做multinomial采样得到展平索引最后把索引还原成(height, width)的坐标probs torch.rand(4, 16, 16) # (batch, height, width) probs_flat probs.flatten(start_dim1) # (4, 256) samples_flat torch.multinomial(probs_flat, num_samples1) # (4, 1) sample_idx samples_flat.squeeze(-1) # (4,) # 把一维索引映射回二维坐标 h_idx sample_idx // 16 w_idx sample_idx % 16 locations torch.stack([h_idx, w_idx], dim1)这个思路就是把多维采样问题拆成“展平采样 坐标还原”两步。如果你有温度参数或者要采样多个位置逻辑也是一样的先用flatten把最后几个维度合并采样后再用除法和取模还原坐标。我实测过这种方法在小 batch 和中等尺寸的张量上不会有任何性能问题代码清晰度反而比一堆permute高得多。记住一个原则multinomial只看最后一个维度所以不管你前面有几维把它 flatten 到让每个样本的概率分量排在最后一维就行。5. 常见问题与排查实录5.1 start_dim 越界与 end_dim 顺序错误flatten最常见的报错有两种。第一种是start_dim或end_dim超出了张量的维度范围比如一个三维张量你传start_dim3PyTorch 会报IndexError或者直接告诉你维度过大。第二种是把start_dim和end_dim写反了比如一个四维张量你写flatten(start_dim2, end_dim1)这会抛出RuntimeError: start_dim must be end_dim。我自己的排查经验是先打印t.shape和t.dim()确认当前张量是几维的再检查参数是否满足0 start_dim end_dim t.dim()。如果你使用了负数索引还要确认转换成正数索引后仍然满足这个条件。比如start_dim-3在四维张量里等价于1如果end_dim-2等价于2这种组合就是合法的。5.2 展平后的形状算不对怎么办很多人会困惑为什么展平后 shape 和自己预想的不一样。根本原因是对“合并维度”的理解有偏差。比如一个形状是(3, 4, 5)的张量flatten(start_dim1)结果是(3, 20)而不是(12, 5)也不是(3, 4, 5)的某种其他变形。20 是 4 和 5 的乘积3 是保留下来的维度大小。我建议你每次做 flatten 之前先做一个“乘积恒等检查”。记录原始每个维度的大小然后把你预期展平后每个维度的大小相乘看看总数是否等于原始总元素数。一旦不等说明你理解错了合并范围。这里有一个具体例子如果你有一个形状(2, 3, 4, 5)的张量flatten(start_dim1, end_dim2)的结果是(2, 12, 5)因为3*412如果你写的是flatten(start_dim2, end_dim3)结果就是(2, 3, 20)。我不止一次看到有人把这两个搞混直接把后面两层带进第一层。5.3 flatten 之后梯度还在吗有个很常见的顾虑flatten 会不会把计算图切断、导致梯度传不回去答案是不会。flatten是一个可微操作它的底层是reshape返回的张量与输入张量共享底层存储连续时自动求导时梯度会沿着原路径回传。你可以放心地在网络 forward 里使用 flatten不需要手动处理任何梯度相关的事情。但有一个细节要注意在 PyTorch 的自动求导体系里view、reshape、flatten这些“形状调整”操作返回的都是原张量的一种视图或复制后的视图它们不改变数据本身只改变解释方式。所以梯度回传时形状会自然对齐。你唯一要小心的是那些原地修改操作比如x.flatten().add_(1)这种操作如果干扰了共享存储的另一个张量可能会在反向传播时引发奇怪的错误。我的建议是flatten 之后的张量不要做原地修改要修改就先用clone()。5.4 快速自查清单每次你在模型里用了 flatten 之后形状不对按下面这个顺序自查五分钟内基本能定位问题打印原始张量的shape确认维度含义尤其是第 0 维是不是 batch。确定你想保留哪些维度、合并哪些维度写出预期结果的数学表达式比如合并后特征长度 dim_1 * dim_2 * dim_3。检查start_dim是否写对最常见的错误是应该start_dim1却写成了0导致 batch 也被合并。检查是否需要end_dim。如果不写默认合并到最后如果你只想合并中间两维必须显式指定。用一个小张量做实验比如torch.randn(2, 3, 4, 5).flatten(...)打印形状对比预期确认后再替换成真实张量。如果用了view报错换成flatten或reshape大概率就好了如果用了flatten还报错检查是不是 start_dim 和 end_dim 的范围问题。这套清单我已经用了很久每次都能快速帮我定位是形状问题、内存连续性问题还是参数问题。6. 我踩过的几个坑6.1 用 view 踩了非连续性坑有一段时间我写代码习惯非常差动不动就x.view(x.size(0), -1)。有一次我对卷积输出做了permute调整通道顺序然后想展平接全连接层结果view直接抛异常。我当时还没反应过来以为是维度算错了排查了半天才发现是内存连续性问题。后来我给自己定了一条规矩凡是可能经过permute、transpose、切片取出的张量一律用flatten绝对不用view。这个小改变让我的 debug 时间缩短了不少。如果你想确认一个张量到底连不连续可以用t.is_contiguous()检查很直观。6.2 把 batch 维度也 flatten 进去这种事情每个 PyTorch 使用者应该都干过至少一次。我刚接触nn.Flatten的时候以为nn.Flatten()和torch.flatten完全一样直接在网络里写了torch.flatten(x)然后送进nn.Linear训练一启动就报维度错误。后来我才意识到nn.Flatten的默认start_dim1是为网络层专门设计的。这个教训告诉我在使用一个 API 之前一定要看默认参数尤其在 PyTorch 这种“同一个概念有多个入口”的框架里函数和类的默认约定可能完全不一样。6.3 一个内存共享引发的反直觉现象还有一个更隐蔽的坑是内存共享。flatten出来的张量和原张量在连续情况下共享底层存储所以如果你对 flatten 后的张量做了原地修改原张量的数据也会变。有一次我在数据处理流水线里写了类似probs_flat[0] 0的代码结果原概率图也被改了导致后面几个样本的采样结果全错了。排查到最后才发现是共享存储的问题。从那以后凡是需要修改 flatten 结果的场景我都会先clone()明确切断存储共享关系。这不是算法上的问题但工程上极其容易踩到分享出来提醒各位。Tensor.flatten(start_dim) 这个接口本身很简单但真正把它用好的关键是你对“维度”这件事的理解深度。一个操作只有在你知道它为什么这样设计、底层做了什么、和 view/reshape 有什么异同之后才能在你的代码里发挥出最大价值。希望这篇内容能帮你少走几次弯路把更多时间留给真正需要动脑的模型设计上。