Tensor.flatten(start_dim)深度解析：PyTorch张量展平与维度控制实战

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_dim

1.1 维度不是玄学：张量的轴到底怎么数

想理解start_dim，第一步不是背参数，而是把“维度”这个词在脑子里变成一张立体的图。你可以把一个张量想象成一个多层书架：第 0 维是书架的层数，第 1 维是每一层上放的书，第 2 维是每本书的页码。比如一个形状是(2, 3, 4)的张量，就是 2 层书架、每层 3 本书、每本书 4 页。

在 PyTorch 里，dim=0永远是最外层，也就是shape元组里的第一个数；dim=1是第二层，以此类推。很多人第一次接触时会把“列”和“行”搞混，因为二维矩阵里dim=0是行、dim=1是列，但一旦上升到三维、四维，这种直觉就不够用了。我自己的习惯是，每拿到一个张量，先写一句注释把每个维度的含义标出来。比如一个图像 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_dim=0，也就是从最外层开始，把所有维度合并成一个一维张量；默认end_dim=-1，也就是一直合并到最后一维。

举个例子，形状为(2, 3, 4)的张量，如果调用t.flatten(start_dim=1)，那第 0 维的 2 保留不动，从第 1 维开始到最后一维合并，结果是(2, 12)。这里 12 就是3 * 4。如果你调用t.flatten(start_dim=0, end_dim=1)，那就是把前两个维度合并，结果是(6, 4)。你看，start_dim就像一个“钉子”，它钉住了前面的维度不让动，只让后面的维度合并起来。

这个参数最常见的用途就是“保留 batch 维度”。在神经网络里，batch是第 0 维，它代表一次喂进去的样本数。当你要把特征图送入全连接层时，不能把 batch 也展平进去，否则全连接层会分不清哪些元素属于哪个样本。这时候写x.flatten(start_dim=1)，就保证了输出还是(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默认是 0，end_dim默认是 -1，所以等价于把 2、3、4 三个维度全部合并成一维。这个操作在把输出特征变成一维向量时很有用，比如你要对某个特征计算距离、做聚类，或者把中间结果送入一个只接受一维输入的处理函数。

这里要注意一个细节：t.flatten()和t.flatten(0)虽然结果相同，但语义上有一点点差别。flatten()用的是默认参数，只要维度变化它都会重新计算范围；flatten(0)是显式指定从第 0 维开始。如果你的张量维度个数是动态变化的，显式写死某个数字并不总是安全。我一般会优先用默认参数或负数索引，让维度的适应性更强。

2.2 用 start_dim 保住 batch 维度

实际写模型时，用到最多的其实是start_dim=1。考虑一个图像分类任务，特征图经过卷积网络后形状是(8, 64, 7, 7)，8 是 batch size，64 是通道数，7x7 是空间尺寸。你想送进全连接层，就需要把通道和空间尺寸合并，但 batch 必须留着。

x = torch.randn(8, 64, 7, 7) y = x.flatten(start_dim=1) 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_dim=1)基本一致，但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_dim=0, end_dim=1) 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_dim=0, end_dim=-1)是一个函数，Tensor.flatten(start_dim=0, end_dim=-1)是一个张量方法，两者等价，只是调用姿势不同。还有一个nn.Flatten(start_dim=1, end_dim=-1)，它是一个网络层模块。

nn.Flatten有一个极其重要的默认值：start_dim=1。它故意把第 0 维也就是 batch 维保留下来，因为绝大多数神经网络模块在使用展平时，都不希望丢失 batch 信息。所以你在nn.Sequential里写nn.Flatten()时，它等价于x.flatten(start_dim=1)，而不是x.flatten(start_dim=0)。这个设计决策让很多新人在把torch.flatten用到网络里时踩坑：你以为是同一个东西，结果 batch 维被吞了。

我建议在写网络层时统一用nn.Flatten()，在写数据处理逻辑时用Tensor.flatten(start_dim=...)，这样语义清晰，也更容易读代码。下面是三者对比：

调用方式	默认 start_dim	默认 end_dim	典型使用场景
torch.flatten(x)	0	-1	把任意张量转成一维向量
x.flatten(start_dim=1)	0	-1	保留 batch，合并特征维度
nn.Flatten()	1	-1	在nn.Sequential中使用

3. flatten、view、reshape 怎么选

3.1 底层机制：view、view+contiguous、reshape 的真实差异

写 PyTorch 的人早晚会遇到view、reshape、flatten这三个操作，它们都能改变张量形状，但底层的逻辑完全不同。view要求张量在内存中是连续的，它只是改变了解释内存的方式，不复制数据。如果张量不连续，比如你对一个多维张量做了transpose，再调用view，PyTorch 会直接抛出RuntimeError: view size is not compatible with input tensor's 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_dim=1) # 正常，结果是 (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_dim=1)一眼就能看出来是保留第 0 维、合并后续。

第三个，你明确知道张量连续，并且追求极致的性能，比如在高频循环里做形状变换，可以用view，省掉 reshape 内部的contiguous检查。但要写注释说明为什么这里安全。

第四，如果你在写一个对输入形状没有严格限制的通用模块，优先flatten(start_dim=1)，它会兼容连续和非连续内存，也几乎不会报错。

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, padding=1), nn.ReLU(), nn.MaxPool2d(2) ) self.flatten = nn.Flatten(start_dim=1) 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_dim=1)和内置的nn.Flatten()效果一模一样，写出来是为了强调 start_dim 在这个场景里的作用。很多人会问：为什么不能用全局平均池化代替展平？可以，但展平保住了所有空间信息，池化会损失一部分。在需要保留细节的任务里，展平是不可替代的。另外，你完全可以在forward里手写x.flatten(start_dim=1)，但把它封装成层的好处是打印模型结构时更清楚。

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_dim=0, end_dim=1) 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_dim=2)，结果是(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_dim=1) # (8, 16*128) = (8, 2048) fused = torch.cat([text_feat, image_feat], dim=1) # (8, 2560)

这里如果你忘了start_dim=1，直接把text_feat展平成一维，就会得到(8*16*128,)，batch 维度消失，拼接必然失败。另外一个非常常见的“送 embedding”场景是：模型输出四维特征图(batch, channels, height, width)，你想得到每个样本的 embedding 向量，可以先做flatten(start_dim=1)得到(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_dim=1) # (4, 256) samples_flat = torch.multinomial(probs_flat, num_samples=1) # (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], dim=1)

这个思路就是把多维采样问题拆成“展平采样 + 坐标还原”两步。如果你有温度参数或者要采样多个位置，逻辑也是一样的：先用flatten把最后几个维度合并，采样后再用除法和取模还原坐标。我实测过，这种方法在小 batch 和中等尺寸的张量上不会有任何性能问题，代码清晰度反而比一堆permute高得多。记住一个原则：multinomial只看最后一个维度，所以不管你前面有几维，把它 flatten 到让每个样本的概率分量排在最后一维就行。

5. 常见问题与排查实录

5.1 start_dim 越界与 end_dim 顺序错误

flatten最常见的报错有两种。第一种是start_dim或end_dim超出了张量的维度范围，比如一个三维张量你传start_dim=3，PyTorch 会报IndexError或者直接告诉你维度过大。第二种是把start_dim和end_dim写反了，比如一个四维张量你写flatten(start_dim=2, end_dim=1)，这会抛出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_dim=1)结果是(3, 20)，而不是(12, 5)，也不是(3, 4, 5)的某种其他变形。20 是 4 和 5 的乘积，3 是保留下来的维度大小。

我建议你每次做 flatten 之前，先做一个“乘积恒等检查”。记录原始每个维度的大小，然后把你预期展平后每个维度的大小相乘，看看总数是否等于原始总元素数。一旦不等，说明你理解错了合并范围。这里有一个具体例子：如果你有一个形状(2, 3, 4, 5)的张量，flatten(start_dim=1, end_dim=2)的结果是(2, 12, 5)，因为3*4=12；如果你写的是flatten(start_dim=2, end_dim=3)，结果就是(2, 3, 20)。我不止一次看到有人把这两个搞混，直接把后面两层带进第一层。

5.3 flatten 之后梯度还在吗

有个很常见的顾虑：flatten 会不会把计算图切断、导致梯度传不回去？答案是不会。flatten是一个可微操作，它的底层是reshape，返回的张量与输入张量共享底层存储（连续时），自动求导时梯度会沿着原路径回传。你可以放心地在网络 forward 里使用 flatten，不需要手动处理任何梯度相关的事情。

但有一个细节要注意：在 PyTorch 的自动求导体系里，view、reshape、flatten这些“形状调整”操作返回的都是原张量的一种视图（或复制后的视图），它们不改变数据本身，只改变解释方式。所以梯度回传时，形状会自然对齐。你唯一要小心的是那些原地修改操作，比如x.flatten().add_(1)，这种操作如果干扰了共享存储的另一个张量，可能会在反向传播时引发奇怪的错误。我的建议是：flatten 之后的张量，不要做原地修改，要修改就先用clone()。

5.4 快速自查清单

每次你在模型里用了 flatten 之后形状不对，按下面这个顺序自查，五分钟内基本能定位问题：

1. 打印原始张量的shape，确认维度含义，尤其是第 0 维是不是 batch。
2. 确定你想保留哪些维度、合并哪些维度，写出预期结果的数学表达式，比如合并后特征长度 = dim_1 * dim_2 * dim_3。
3. 检查start_dim是否写对，最常见的错误是应该start_dim=1却写成了0，导致 batch 也被合并。
4. 检查是否需要end_dim。如果不写，默认合并到最后；如果你只想合并中间两维，必须显式指定。
5. 用一个小张量做实验，比如torch.randn(2, 3, 4, 5).flatten(...)，打印形状对比预期，确认后再替换成真实张量。
6. 如果用了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_dim=1是为网络层专门设计的。这个教训告诉我：在使用一个 API 之前，一定要看默认参数，尤其在 PyTorch 这种“同一个概念有多个入口”的框架里，函数和类的默认约定可能完全不一样。

6.3 一个内存共享引发的反直觉现象

还有一个更隐蔽的坑是内存共享。flatten出来的张量和原张量在连续情况下共享底层存储，所以如果你对 flatten 后的张量做了原地修改，原张量的数据也会变。有一次我在数据处理流水线里写了类似probs_flat[0] = 0的代码，结果原概率图也被改了，导致后面几个样本的采样结果全错了。排查到最后才发现是共享存储的问题。从那以后，凡是需要修改 flatten 结果的场景，我都会先clone()，明确切断存储共享关系。这不是算法上的问题，但工程上极其容易踩到，分享出来提醒各位。

Tensor.flatten(start_dim) 这个接口本身很简单，但真正把它用好的关键，是你对“维度”这件事的理解深度。一个操作只有在你知道它为什么这样设计、底层做了什么、和 view/reshape 有什么异同之后，才能在你的代码里发挥出最大价值。希望这篇内容能帮你少走几次弯路，把更多时间留给真正需要动脑的模型设计上。