PyTorch 常见层函数：池化、卷积、线性、归一化、Embedding

Abstract

这篇按“层函数类别”组织，而不是按模型组织。

每个函数都回答三件事：它是什么、一个最小例子、它在 DLinear / Informer / PatchTST 的哪个源码位置出现。

0. 文件索引

项目	内容
本文主题	PyTorch 常见 nn.* 层函数
覆盖范围	池化层、卷积层、线性层、归一化层、正则化层、padding 层、embedding 层、输出头
主要模型	DLinear / Informer / PatchTST
配套文档	[[11-PyTorch-Tensor基础操作-切片变形拼接注意力]]

0.1 知识地图

类别	函数	最重要的 shape 规则	出现位置
池化层	nn.AvgPool1d	输入 (B,C,L)	DLinear moving_avg
池化层	nn.MaxPool1d	输入 (B,C,L)，常用于下采样	Informer ConvLayer
卷积层	nn.Conv1d	输入 (B,C,L)	Informer embedding / Transformer FFN
线性层	nn.Linear	只作用最后一维	DLinear / Informer / PatchTST 都有
归一化层	nn.LayerNorm	通常归一化最后一维	Informer / PatchTST Encoder
归一化层	nn.BatchNorm1d	通常处理 (B,C,L)	Informer ConvLayer
正则化层	nn.Dropout	不改 shape	embedding / attention / FFN / head
padding 层	nn.ReplicationPad1d	输入 (B,C,L)，复制边界	PatchTST patching
embedding 层	nn.Embedding	离散 id -> 向量	Informer 时间特征
输出头	nn.Flatten	把多个维度展平	PatchTST FlattenHead
容器/参数	nn.ModuleList / nn.Parameter	管理子层 / 显式参数	DLinear individual 分支

---

0.2 数学公式总览：先看“主导维度”

大多数层函数的抽象逻辑可以先分成 5 类：

类别	主导逻辑	典型层函数	一句话理解
滑窗聚合	沿长度维 L 滑动窗口	AvgPool1d / MaxPool1d / Conv1d(kernel>1)	用窗口看局部片段
最后一维投影	对最后一维做矩阵乘法	Linear	(..., in) -> (..., out)
统计归一化	用均值方差标准化	LayerNorm / BatchNorm1d	关键是统计哪些维度
随机正则	训练时随机置零	Dropout	不改 shape，只改数值
查表/变形	查表或合并维度	Embedding / Flatten	一个是取参数行，一个是改 shape

![[zdocs/pytorch-basics/assets/concept_layers_formula_map.svg]]

---

1. 池化层：AvgPool1d 和 MaxPool1d

1.1 nn.AvgPool1d

基本作用：

沿一维序列滑动窗口，取窗口内平均值。

输入格式：

(B, C, L)

主导公式：

$$\begin{array}{rl}{y}_{b,c,i}& =\mathrm{mean}\left(x_{b,c,star{t}_i:star{t}_i+k}\right)\\ star{t}_i& =i\cdot s\end{array}$$

其中 k = kernel_size，s = stride。

更一般地，输出长度由下面公式决定：

$$L_{out}=\lfloor \frac{L_{in}+2p-d(k-1)-1}{s}+1\rfloor$$

其中 p = padding，d = dilation。

AvgPool1d 一般不用 dilation，所以可简化为：

$$L_{out}=\lfloor \frac{L_{in}+2p-k}{s}+1\rfloor$$

![[zdocs/pytorch-basics/assets/concept_window_length_formula.svg]]

最小例子：

import torch
import torch.nn as nn

x = torch.tensor([[[1., 1., 2., 3., 4., 5., 6., 6.]]])
# shape: (B=1, C=1, L=8)

pool = nn.AvgPool1d(kernel_size=3, stride=1, padding=0)
y = pool(x)

print(y)

输出对应窗口：

[1, 1, 2] -> 1.333
[1, 2, 3] -> 2.000
[2, 3, 4] -> 3.000
[3, 4, 5] -> 4.000
[4, 5, 6] -> 5.000
[5, 6, 6] -> 5.667

shape：

(1, 1, 8) -> (1, 1, 6)

这里：

$$\begin{array}{rl}{L}_{in}& =8,k=3,s=1,p=0\\ L_{out}& =\lfloor \frac{8+2\cdot 0-3}{1}+1\rfloor =6\end{array}$$

在我们模型里的位置：

模型	源码位置	作用
DLinear	Autoformer_EncDec.py -> moving_avg.__init__	对时间序列做滑动平均，得到 trend

源码：

self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)

forward 中：

x = self.avg(x.permute(0, 2, 1))
x = x.permute(0, 2, 1)

为什么要 permute：

DLinear 主线: (B,T,C)
AvgPool1d 要求: (B,C,L)

所以：

(2,8,3) -> permute(0,2,1) -> (2,3,8) -> AvgPool1d -> (2,3,6) -> permute -> (2,6,3)

详细下钻：[[../model-order/01-DLinear/02-AvgPool1d与permute-DLinear-moving_avg|02-AvgPool1d与permute-DLinear-moving_avg]]

1.2 nn.MaxPool1d

基本作用：

沿一维序列滑动窗口，取窗口内最大值，常用于下采样。

主导公式：

$$\begin{array}{rl}{y}_{b,c,i}& =max\left(x_{b,c,star{t}_i:star{t}_i+k}\right)\\ star{t}_i& =i\cdot s\end{array}$$

输出长度和 AvgPool1d 同一套公式：

$$L_{out}=\lfloor \frac{L_{in}+2p-d(k-1)-1}{s}+1\rfloor$$

最小例子：

x = torch.tensor([[[1., 3., 2., 5., 4., 6.]]])
pool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
y = pool(x)

输出长度：

$$\begin{array}{rl}{L}_{out}& =\lfloor \frac{L_{in}+2p-d(k-1)-1}{s}+1\rfloor \\ & =\lfloor \frac{6+2\cdot 1-1\cdot (3-1)-1}{2}+1\rfloor \\ & =3\end{array}$$

在我们模型里的位置：

模型	源码位置	作用
Informer	Transformer_EncDec.py -> ConvLayer.maxPool	encoder distilling，把序列长度压短

源码：

self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)

---

2. 卷积层：nn.Conv1d

基本作用：

用可学习的局部窗口沿一维序列滑动，提取局部模式。

输入格式：

(B, C, L)

主导公式：

$$y_{b,c_{out},i}=b_{c_{out}}+\sum _{c_{in}}\sum _{r=0}^{k-1}{W}_{c_{out},c_{in},r}\cdot x_{b,c_{in},i\cdot s+r\cdot d-p}$$

理解重点：

AvgPool1d: 窗口权重固定，全部取平均
MaxPool1d: 窗口内只取最大值
Conv1d: 窗口权重可学习，并且可以把 C_in 映射到 C_out

输出长度仍然是滑窗公式：

$$L_{out}=\lfloor \frac{L_{in}+2p-d(k-1)-1}{s}+1\rfloor$$

2.1 kernel_size=3：看局部时间窗口

最小例子：

conv = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=3, bias=False)

如果权重是：

w = [0.25, 0.50, 0.25]

输入：

x = [2, 4, 6, 8, 10]

第一个窗口输出：

2*0.25 + 4*0.50 + 6*0.25 = 4

如果 stride=1、padding=0，则：

$$\begin{array}{rl}{L}_{in}& =5,k=3,s=1,p=0\\ L_{out}& =\lfloor \frac{5-3}{1}+1\rfloor =3\end{array}$$

在我们模型里的位置：

模型	源码位置	作用
Informer	Embed.py -> TokenEmbedding.tokenConv	把原始变量值编码到 d_model
Informer	Transformer_EncDec.py -> ConvLayer.downConv	distilling 前的局部卷积

2.2 kernel_size=1：position-wise FFN

kernel_size=1 不看左右邻居，只对当前位置的 channel 维做投影。

此时卷积公式退化成每个位置独立的线性变换：

$$y_{b,c_{out},i}=b_{c_{out}}+\sum _{c_{in}}{W}_{c_{out},c_{in},0}\cdot x_{b,c_{in},i}$$

它不会混合不同位置 i，只混合 channel 维。

在 Transformer 里常用它实现 FFN：

self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)

shape：

(B,L,d_model)
-> transpose
(B,d_model,L)
-> Conv1d(d_model -> d_ff, kernel=1)
(B,d_ff,L)
-> Conv1d(d_ff -> d_model, kernel=1)
(B,d_model,L)
-> transpose
(B,L,d_model)

在我们模型里的位置：

模型	源码位置	作用
Informer	Transformer_EncDec.py -> EncoderLayer / DecoderLayer	FFN
PatchTST	Transformer_EncDec.py -> EncoderLayer	FFN

详细下钻：[[../model-order/02-PatchTST/03-Conv1d与BCL格式-Informer-PatchTST|03-Conv1d与BCL格式-Informer-PatchTST]]

---

3. 线性层：nn.Linear

基本规则：

nn.Linear(in_features, out_features) 只作用在输入 tensor 的最后一维。

统一 shape：

(..., in_features) -> (..., out_features)

主导公式：

$$y_{\dots ,j}=b_j+\sum _i{x}_{\dots ,i}{W}_{j,i}$$

矩阵写法：

$$y=x{W}^{\mathrm{\top }}+b$$$$W\in {\mathbb{R}}^{out\mathrm{\_}features\times in\mathrm{\_}features},b\in {\mathbb{R}}^{out\mathrm{\_}features}$$

最小例子：

linear = nn.Linear(3, 2)
x = torch.randn(2, 6, 3)
y = linear(x)

shape：

(2,6,3) -> (2,6,2)

在我们模型里的位置：

模型	源码位置	作用
DLinear	DLinear.py -> Linear_Seasonal / Linear_Trend	seq_len -> pred_len
Informer	SelfAttention_Family.py -> AttentionLayer	Q/K/V 投影、输出投影
Informer	Informer.py -> projection	d_model -> c_out
PatchTST	Embed.py -> PatchEmbedding.value_embedding	patch_len -> d_model
PatchTST	PatchTST.py -> FlattenHead.linear	head_nf -> pred_len

DLinear 特别容易混：

原始: (B,T,C)
先 permute: (B,C,T)
Linear(seq_len, pred_len)
输出: (B,C,pred_len)

详细下钻：[[../model-order/01-DLinear/03-Linear最后一维规则-DLinear-Informer-PatchTST|03-Linear最后一维规则-DLinear-Informer-PatchTST]]

---

4. 归一化层：LayerNorm 和 BatchNorm1d

两者共享同一个抽象公式：

$$y=\gamma \cdot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$

区别不是公式，而是 mean 和 var 到底沿哪些维度统计。

![[zdocs/pytorch-basics/assets/concept_norm_dimension_compare.svg]]

4.1 nn.LayerNorm

常见输入：

(B, L, d_model)

作用：

对每个 token 的最后一维 d_model 做归一化。

对输入 x.shape=(B,L,d_model)：

$${\mu }_{b,l}=\frac{1}{d_{model}}\sum _i{x}_{b,l,i}$$$${\sigma }_{b,l}^2=\frac{1}{d_{model}}\sum _i{(x_{b,l,i}-{\mu }_{b,l})}^2$$$$y_{b,l,i}={\gamma }_i\cdot \frac{x_{b,l,i}-{\mu }_{b,l}}{\sqrt{{\sigma }_{b,l}^2+ϵ}}+{\beta }_i$$

也就是说：

每个 token 自己算自己的 mean/var。
不跨 batch，也不跨时间位置。

最小例子：

norm = nn.LayerNorm(16)
x = torch.randn(2, 6, 16)
y = norm(x)

shape：

(2,6,16) -> (2,6,16)

在我们模型里的位置：

模型	源码位置	作用
Informer	Transformer_EncDec.py -> EncoderLayer / DecoderLayer	残差后归一化
PatchTST	Transformer_EncDec.py -> EncoderLayer	残差后归一化

4.2 nn.BatchNorm1d

常见输入：

(B, C, L)

作用：

对 channel 维做 batch/length 统计归一化。

对输入 x.shape=(B,C,L)：

$${\mu }_c=\frac{1}{B\cdot L}\sum _b\sum _l{x}_{b,c,l}$$$${\sigma }_c^2=\frac{1}{B\cdot L}\sum _b\sum _l{(x_{b,c,l}-{\mu }_c)}^2$$$$y_{b,c,l}={\gamma }_c\cdot \frac{x_{b,c,l}-{\mu }_c}{\sqrt{{\sigma }_c^2+ϵ}}+{\beta }_c$$

也就是说：

每个 channel 单独统计；
统计范围覆盖 batch 维 B 和长度维 L。

在我们模型里的位置：

模型	源码位置	作用
Informer	Transformer_EncDec.py -> ConvLayer.norm	Conv1d 后稳定 channel 分布

详细下钻：[[../model-order/02-PatchTST/04-LayerNorm-BatchNorm-Dropout-Transformer基础|04-LayerNorm-BatchNorm-Dropout-Transformer基础]]

---

5. 正则化层：nn.Dropout

基本作用：

训练时随机把部分元素置 0，推理时不随机置零。

训练时主导公式：

$$m\sim \mathrm{Bernoulli}(1-p)$$$$y=\frac{x\cdot m}{1-p}$$

其中：

p: 被置零的概率
1-p: 被保留的概率

PyTorch 训练时会除以 1-p，是为了让输出期望保持不变：

$$\mathbb{E}[y]=x$$

最小例子：

dropout = nn.Dropout(p=0.5)
x = torch.ones(4)
y = dropout(x)

训练时可能得到：

[0, 2, 0, 2]

shape 不变。

在我们模型里的位置：

模型	源码位置	作用
Informer	Embed.py / SelfAttention_Family.py / Transformer_EncDec.py	embedding、attention、FFN 正则化
PatchTST	Embed.py / PatchTST.py	patch embedding、head 正则化
DLinear	DLinear.py classification	分类分支正则化

---

6. Padding 层：nn.ReplicationPad1d

基本作用：

沿最后一维复制边界值做 padding。

输入格式：

(B, C, L)

长度公式：

$$L_{out}=L_{in}+p_{left}+p_{right}$$

数值规则：

左侧补的值 = 原序列第一个值
右侧补的值 = 原序列最后一个值

最小例子：

pad = nn.ReplicationPad1d((0, 2))
x = torch.tensor([[[1., 2., 3.]]])
y = pad(x)

输出：

[1, 2, 3] -> [1, 2, 3, 3, 3]

在我们模型里的位置：

模型	源码位置	作用
PatchTST	Embed.py -> PatchEmbedding.padding_patch_layer	右端复制，让最后一个 patch 可切出完整长度

详细下钻：[[../model-order/02-PatchTST/01-ReplicationPad1d与unfold-PatchTST-PatchEmbedding|01-ReplicationPad1d与unfold-PatchTST-PatchEmbedding]]

---

7. Embedding 层：nn.Embedding 与自定义位置编码

7.1 nn.Embedding

基本作用：

把离散 id 查表成向量。

主导公式：

$$y=W[id]$$$$W\in {\mathbb{R}}^{num\mathrm{\_}embeddings\times embedding\mathrm{\_}dim}$$

如果输入 id 的 shape 是：

id.shape = (...)

则输出 shape 是：

$$y.shape=(...,embedding\mathrm{\_}dim)$$

最小例子：

emb = nn.Embedding(num_embeddings=24, embedding_dim=16)
hour_id = torch.tensor([13])
hour_vec = emb(hour_id)

shape：

(1,) -> (1,16)

在我们模型里的位置：

模型	源码位置	作用
Informer	Embed.py -> TemporalEmbedding	month/day/weekday/hour/minute 时间特征 embedding

7.2 PositionalEmbedding

基本作用：

用正余弦函数生成位置编码，让 attention 知道 token 顺序。

主导公式：

$$P{E}_{pos,2i}=\sin \left(\frac{pos}{{10000}^{2i/d_{model}}}\right)$$$$P{E}_{pos,2i+1}=\cos \left(\frac{pos}{{10000}^{2i/d_{model}}}\right)$$

它不是根据输入值大小生成的，而是根据 token 的位置 pos 生成的。

在我们模型里的位置：

模型	源码位置	作用
Informer	Embed.py -> DataEmbedding.position_embedding	时间步位置编码
PatchTST	Embed.py -> PatchEmbedding.position_embedding	patch 位置编码

详细下钻：[[../model-order/03-Informer/01-Embedding与位置编码-Informer-PatchTST|01-Embedding与位置编码-Informer-PatchTST]]

---

8. 输出头：nn.Flatten

基本作用：

从指定维度开始，把多个维度合并成一个维度。

主导公式：

$$new\mathrm{\_}dim=\prod _{j=start}^{end}di{m}_j$$

更准确地说，Flatten 只改变 shape，不改变元素顺序，不学习参数。

最小例子：

flatten = nn.Flatten(start_dim=-2)
x = torch.randn(2, 4, 16, 6)
y = flatten(x)

shape：

(2,4,16,6) -> (2,4,96)

这里：

$$96=16\times 6$$

在我们模型里的位置：

模型	源码位置	作用
PatchTST	PatchTST.py -> FlattenHead	把 d_model * patch_num 展平成预测头输入

详细下钻：[[../model-order/02-PatchTST/02-Flatten与标准化统计-PatchTST输出头|02-Flatten与标准化统计-PatchTST输出头]]

---

9. 容器与参数：ModuleList 和 Parameter

9.1 nn.ModuleList

基本作用：

保存一组子模块，让 PyTorch 能正确注册它们的参数。

在我们模型里的位置：

模型	源码位置	作用
DLinear	DLinear.py -> individual=True	每个变量一个独立 Linear
Informer / PatchTST	Transformer_EncDec.py -> Encoder / Decoder	保存多层 EncoderLayer / DecoderLayer

9.2 nn.Parameter

基本作用：

把一个 tensor 显式注册成可训练参数。

它本身不定义新的数学变换，只决定这个 tensor 是否被优化器更新：

普通 tensor: optimizer 不更新
nn.Parameter: optimizer 会更新

在 DLinear 里：

self.Linear_Seasonal.weight = nn.Parameter(
    (1 / self.seq_len) * torch.ones([self.pred_len, self.seq_len])
)

含义：

手动把 Linear 权重初始化成平均权重。

---

10. 一句话总结

这一组 nn.* 层函数可以这样记：

AvgPool1d / MaxPool1d: 沿时间轴滑窗聚合
Conv1d: 输入必须是 (B,C,L)，可看局部窗口或做 FFN
Linear: 只改最后一维
LayerNorm: 常看最后一维
BatchNorm1d: 常看 channel 维
Dropout: 训练时随机置零，不改 shape
ReplicationPad1d: 复制边界补齐
Embedding: 离散 id 或位置 -> d_model 向量
Flatten: 多维展平给输出头