深度学习入门到入神·悟道篇01：Tensor 变形术——Flatten、GAP、mean、keepdim 到底差在哪？

从入门到入神｜PyTorch 基础番外

这篇文章建议大家看完我们的深度学习入门第四关之后再看，这样会对这个内容理解的更加的深刻。

学到 Flatten 和 Global Average Pooling 以后，代码里经常会同时出现：

torch.flatten(x, 1)

nn.AdaptiveAvgPool2d((1, 1))

x.mean(dim=(2, 3))

x.mean(dim=(2, 3), keepdim=True)

它们看起来都能让 Tensor 的尺寸变小，但做的事情其实完全不同。

这一篇不增加新的网络知识，只解决一个问题：

这几个操作到底有什么区别？

假设我们始终使用：

x = torch.randn(8, 128, 8, 8)

也就是：

x.shape = [8, 128, 8, 8]

             ↑    ↑   ↑  ↑
             B    C   H  W

---

一、先看完整代码

先把四种情况全部跑一遍：

import torch
import torch.nn as nn


x = torch.randn(8, 128, 8, 8)

print("原始 Tensor：")
print(x.shape)


# 1. Flatten
x_flatten = torch.flatten(x, 1)

print("\n1. flatten(x, 1)：")
print(x_flatten.shape)


# 2. Global Average Pooling
gap = nn.AdaptiveAvgPool2d((1, 1))

x_gap = gap(x)

print("\n2. AdaptiveAvgPool2d((1,1))：")
print(x_gap.shape)


# 3. mean
x_mean = x.mean(dim=(2, 3))

print("\n3. mean(dim=(2,3))：")
print(x_mean.shape)


# 4. mean + keepdim
x_mean_keepdim = x.mean(
    dim=(2, 3),
    keepdim=True
)

print("\n4. mean(dim=(2,3), keepdim=True)：")
print(x_mean_keepdim.shape)

运行结果：

原始 Tensor：
torch.Size([8, 128, 8, 8])

1. flatten(x, 1)：
torch.Size([8, 8192])

2. AdaptiveAvgPool2d((1,1))：
torch.Size([8, 128, 1, 1])

3. mean(dim=(2,3))：
torch.Size([8, 128])

4. mean(dim=(2,3), keepdim=True)：
torch.Size([8, 128, 1, 1])

先把结果放在一起：

操作	输出 Shape
原始 Tensor	[8,128,8,8]
flatten(x,1)	[8,8192]
AdaptiveAvgPool2d((1,1))	[8,128,1,1]
mean(dim=(2,3))	[8,128]
mean(dim=(2,3), keepdim=True)	[8,128,1,1]

接下来逐个看。

---

二、torch.flatten(x, 1)：只是展开，不做平均

输入：

[8,128,8,8]

执行：

x = torch.flatten(x, 1)

这里的 1 表示：

从第 1 维开始展开。

原来的维度编号是：

[8, 128, 8, 8]
 ↑   ↑   ↑  ↑
 0   1   2  3

所以：

第0维 Batch 保留

第1、2、3维合并

于是：

128 × 8 × 8
=
8192

最终：

[8,128,8,8]
↓
flatten(x,1)
↓
[8,8192]

关键点：

Flatten 没有求平均，也没有删除这些数，只是重新组织 Shape。

---

三、AdaptiveAvgPool2d((1,1))：每个 Channel 压成 1×1

执行：

gap = nn.AdaptiveAvgPool2d((1, 1))

y = gap(x)

输入：

[8,128,8,8]

输出：

[8,128,1,1]

为什么？

因为 (1,1) 表示：

每张 Feature Map 最终都压缩成 1×1。

对于一个 Channel：

8×8
↓
64个位置求平均
↓
1×1

128个 Channel 分别这样做：

128张 8×8 Feature Map
↓
128张 1×1 Feature Map

所以：

[8,128,8,8]
↓
GAP
↓
[8,128,1,1]

注意：

GAP 会做平均，因此它和 Flatten 本质不同。

---

四、mean(dim=(2,3))：直接对 H、W 求平均

再来看：

y = x.mean(dim=(2, 3))

为什么是：

dim=(2,3)

因为：

[8,128,8,8]
 ↑  ↑   ↑ ↑
 0  1   2 3

0 = Batch
1 = Channel
2 = Height
3 = Width

所以：

x.mean(dim=(2, 3))

就是：

对 H 和 W 两个维度求平均。

于是：

[8,128,8,8]
↓
H、W求平均
↓
[8,128]

这里它实际上完成的计算和 GAP 的核心思想非常接近：

每个 Channel 最后得到一个平均值。

但是输出 Shape 有区别。

---

五、keepdim=True 到底在 keep 什么？

这是最容易让初学者迷糊的一点。

普通：

x.mean(dim=(2, 3))

输出：

[8,128]

因为 H、W 两个维度求平均以后直接消失了。

如果写：

x.mean(
    dim=(2, 3),
    keepdim=True
)

输出：

[8,128,1,1]

区别就是：

keepdim=False（默认）

[8,128,8,8]
↓
[8,128]

H、W 直接消失。

keepdim=True

[8,128,8,8]
↓
[8,128,1,1]

H、W 还在，只不过：

H = 1
W = 1

所以：

keepdim = keep dimension = 保留维度。

---

六、为什么 [8,128,1,1] 和 [8,128] 看起来很像，但又不一样？

这两个 Tensor：

[8,128,1,1]

和：

[8,128]

包含的元素数量是一样的：

8 × 128 × 1 × 1
=
8 × 128

但是：

Shape 不一样。

第一个是四维 Tensor：

[B,C,H,W]

第二个是二维 Tensor：

[B,Features]

例如：

[8,128,1,1]
↓
flatten(x,1)
↓
[8,128]

所以可以理解成：

数据数量基本一样，但组织形式不同。

这也是为什么有些网络层希望得到：

[B,C,1,1]

而 Linear 前通常更常见：

[B,C]

---

七、AdaptiveAvgPool2d((1,1)) 和 mean(dim=(2,3), keepdim=True) 是不是一样？

对于我们现在这个 GAP 场景，可以先这样理解：

nn.AdaptiveAvgPool2d((1, 1))

得到：

[8,128,1,1]

而：

x.mean(
    dim=(2,3),
    keepdim=True
)

也得到：

[8,128,1,1]

二者都是：

每个 Channel 对整个 H、W 求平均，并保留 1×1 的空间维度。

所以在这个特定场景下，它们实现的结果可以非常接近。

但是代码表达方式不同：

AdaptiveAvgPool2d((1,1))

强调的是：

我要一个 1×1 的池化输出。

而：

mean(dim=(2,3))

强调的是：

我要对第 2、3 维求平均。

初学阶段先理解这个区别就够了。

---

八、flatten(x, 0)、flatten(x, 1)、flatten(x, 4) 又分别是什么意思？

这里顺便把 Flatten 的数字彻底讲清。

假设：

x.shape = [8,128,8,8]

维度编号：

[8,128,8,8]
 ↑  ↑   ↑ ↑
 0  1   2 3

flatten(x, 1)

torch.flatten(x, 1)

表示：

从第1维开始展开。

所以：

[8,128,8,8]
↓
[8,8192]

Batch 保留。

---

flatten(x, 0)

torch.flatten(x, 0)

表示：

从第0维开始全部展开。

所以：

[8,128,8,8]
↓
[65536]

连 Batch 都被拍平了。

训练网络时通常不是我们想要的结果。

---

flatten(x, 4)

对于：

[8,128,8,8]

只有：

dim=0、1、2、3

不存在：

dim=4

所以：

torch.flatten(x, 4)

会报错。

一定要记住：

维度编号从 0 开始。

---

九、一张表彻底区分这几个操作

写法	做了什么	是否求平均	输出
flatten(x,1)	C、H、W 展开	否	[8,8192]
AdaptiveAvgPool2d((1,1))	每个 Channel 池化成1×1	是	[8,128,1,1]
mean(dim=(2,3))	H、W求平均	是	[8,128]
mean(dim=(2,3), keepdim=True)	H、W求平均并保留维度	是	[8,128,1,1]

再把它画成数据流：

原始：
[8,128,8,8]

Flatten

[8,128,8,8]
↓
[8,8192]

GAP

[8,128,8,8]
↓
[8,128,1,1]

mean

[8,128,8,8]
↓
[8,128]

mean + keepdim

[8,128,8,8]
↓
[8,128,1,1]

---

总结

这一篇最需要记住的不是代码，而是每个操作究竟在干什么。

Flatten

重新整理 Shape，不求平均。

GAP

每个 Channel 对整个 H、W 求平均，输出 1×1。

mean(dim=(2,3))

直接告诉 PyTorch：我要对 H、W 求平均。

keepdim=True

平均以后不要把 H、W 维度删除，让它们保留为 1×1。

而：

torch.flatten(x, 1)

里面的：

1

也不是“变成一维”的意思。

它真正表示：

从第1维开始展开。

理解了这些以后，再看到：

[8,128,8,8]
[8,128,1,1]
[8,128]
[8,8192]

就不会觉得这些 Shape 是凭空变化的了。

它们背后都对应着一个非常具体的 Tensor 操作。

从入门，到入神。