pytorch学习笔记1

从 nn.Linear 到 Conv2d 再到 Attention：PyTorch 代码参数与 Tensor Shape 扫盲笔记

最近在补 D2L 和 PyTorch，为了后面阅读，发现真正经常卡住我的并不是公式本身，而是代码中的一些细节：

• shape[1] 到底是什么意思？
• nn.Linear(128, 64) 谁是输入、谁是输出？
• unsqueeze(-1) 为什么突然多了一维？
• optimizer.step() 到底做了什么？
• Conv2d 中的 Cin、Cout 到底在哪里？
• repeat_interleave 为什么要复制 query？
• torch.bmm 和普通矩阵乘法有什么区别？

所以把目前遇到的 PyTorch 基础 API、参数、Tensor Shape 和神经网络代码逻辑统一整理一下。

---

1. Tensor 创建：先看"作者想造什么数据"

常见的 Tensor 创建函数：

API	含义	常见参数
torch.rand(shape)	生成 (U(0,1)) 均匀随机数	shape
torch.randn(shape)	生成 (N(0,1)) 标准正态随机数	shape
torch.normal(mean, std, size)	生成指定均值、标准差的正态随机数	mean、std、size
torch.zeros(shape)	生成全 0 Tensor	shape
torch.ones(shape)	生成全 1 Tensor	shape
torch.arange(start, stop, step)	生成等间距数列	start、stop、step
torch.tensor(...)	根据已有数据创建 Tensor	data
torch.zeros_like(X)	创建与 X 同 shape 的全 0 Tensor	X

例如：

torch.normal(0.0, 0.5, (50,))

不要把它背成"第一个参数是什么、第二个参数是什么"。

应该直接翻译成：

我要生成 50 个正态随机噪声
均值 = 0
标准差 = 0.5
shape = [50]

也就是：

ϵi∼N(0,0.52) \epsilon_i \sim N(0, 0.5^2) ϵi∼N(0,0.52)

1.1 (1,) 为什么有一个逗号？

Python 中：

(1)

实际上仍然只是整数：

1

而：

(1,)

才表示一个长度为 1 的 tuple。

因此：

torch.rand((1,))

表示生成 shape 为：

[1]

的 Tensor。

实际上也可以直接写：

torch.rand(1)

---

2. torch.sort、索引和 shape

values, indices = torch.sort(X)

torch.sort() 会返回两个结果：

values  = 排序后的值
indices = 排序后的元素原来所在的位置

例如：

x = torch.tensor([4., 1., 3.])
values, indices = torch.sort(x)

结果相当于：

values  = [1., 3., 4.]
indices = [1, 2, 0]

因此：

x_train, _ = torch.sort(...)

中的：

_

表示：

第二个返回值我不使用。

---

2.1 shape[1] 不是取第 1 列

假设：

X.shape == [50, 49]

那么：

X.shape[0]   # 50
X.shape[1]   # 49

表示：

第 0 个维度长度为 50

第 1 个维度长度为 49

而：

X[1]

才是真正对 Tensor 做索引。

所以一定区分：

X.shape[1]

和：

X[1]

---

3. reshape、Flatten、unsqueeze

这些操作贯穿整个深度学习代码。

---

3.1 reshape

X.reshape(new_shape)

作用：

改变 Tensor 的组织方式，但不改变元素总数。

例如：

[2, 3, 4]

总元素数：

2×3×4=24 2\times3\times4=24 2×3×4=24

可以 reshape 为：

[6, 4]

因为：

6×4=24 6\times4=24 6×4=24

---

3.2 -1 是自动推断维度

例如：

X.reshape(-1, 784)

意思是：

第二维我指定为 784，第一维请 PyTorch 根据元素总数自动推断。

---

3.3 nn.Flatten()

Fashion-MNIST 的一张图片：

28×28=784 28\times28=784 28×28=784

因此：

nn.Flatten()

可以把：

[B, 1, 28, 28]

变成：

[B, 784]

然后送入：

nn.Linear(784, 10)

---

3.4 unsqueeze

unsqueeze 的作用：

在指定位置插入一个长度为 1 的维度。

假设：

X.shape = [2, 3]

执行：

X.unsqueeze(1)

得到：

[2, 1, 3]

执行：

X.unsqueeze(-1)

得到：

[2, 3, 1]

其中：

-1

表示最后一维。

Attention 中经常利用它进行矩阵乘法，例如：

[Q, K]
↓ unsqueeze(1)
[Q, 1, K]

---

4. nn.Linear(in_features, out_features)

参数顺序：

nn.Linear(Din, Dout)

表示：

Din→Dout D_{in}\rightarrow D_{out} Din→Dout

例如：

nn.Linear(128, 64)

就是：

128→64 128\rightarrow64 128→64

输入：

[B, 128]

输出：

[B, 64]

---

4.1 为什么 weight.shape 又是反过来的？

对于：

layer = nn.Linear(128, 64)

PyTorch 内部：

layer.weight.shape

是：

[64, 128]

即：

Dout,Din\] \[D_{out}, D_{in}\] \[Dout,Din

因为内部计算等价于：

Y=XWT+b Y=XW^T+b Y=XWT+b

所以需要记住：

Linear 参数顺序：
(in_features, out_features)

weight.shape：
[out_features, in_features]

---

4.2 Linear 参数量

nn.Linear(Din, Dout)

权重数量：

DinDout D_{in}D_{out} DinDout

bias 数量：

Dout D_{out} Dout

总参数：

DinDout+Dout \boxed{D_{in}D_{out}+D_{out}} DinDout+Dout

原因是：

每一个输出神经元都需要读取全部 DinD_{in}Din 个输入，并拥有一个自己的 bias。

---

5. nn.Sequential

例如：

net = nn.Sequential(
    nn.Linear(20, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

可以直接读成：

20
↓ Linear
256
↓ ReLU
256
↓ Linear
10

如果 batch size 为 (B)：

[B, 20]
→ [B, 256]
→ [B, 256]
→ [B, 10]

Sequential 本质就是：

按顺序把上一层的输出作为下一层输入。

因此：

net(X)

会自动依次执行所有模块。

---

6. nn.Module、class、self、__init__、forward

自己定义网络时：

class Net(nn.Module):

表示：

定义一个名为 Net 的 PyTorch 网络模块。

---

6.1 __init__

def __init__(self):

负责定义：

网络中有哪些层、参数和组件。

例如：

self.fc1 = nn.Linear(20, 256)
self.fc2 = nn.Linear(256, 10)

---

6.2 self

self 可以理解成：

当前创建出来的这个网络对象本身。

所以：

self.fc1

就是：

当前这个 Net 实例自己的 fc1 层。

---

6.3 super().__init__()

super().__init__()

调用父类 nn.Module 的初始化逻辑。

这样 PyTorch 才可以自动管理：

parameters
梯度
CPU/GPU
train/eval 模式
子模块

---

6.4 forward

def forward(self, X):

定义：

输入数据进入网络后，具体如何计算输出。

例如：

def forward(self, X):
    X = self.fc1(X)
    X = F.relu(X)
    X = self.fc2(X)
    return X

就是前向传播。

---

7. Forward 和 Backward

Forward

前向传播：

X→Y^→L X\rightarrow\hat Y\rightarrow L X→Y^→L

即：

使用当前参数计算模型预测，再计算 loss。

---

Backward

loss.backward()

不是：

把 (f(x)) 反解回 (x)。

而是：

从 loss 出发，通过链式法则反向计算每个参数的梯度。

例如：

∂L∂W3→∂L∂W2→∂L∂W1 \frac{\partial L}{\partial W_3} \rightarrow \frac{\partial L}{\partial W_2} \rightarrow \frac{\partial L}{\partial W_1} ∂W3∂L→∂W2∂L→∂W1∂L

最后梯度存放在：

parameter.grad

中。

---

8. SGD 到底做什么？

手写 SGD：

def sgd(params, lr, batch_size):
    with torch.no_grad():
        for param in params:
            param -= lr * param.grad / batch_size
            param.grad.zero_()

核心公式：

θ←θ−η∇θL \theta \leftarrow \theta-\eta\nabla_\theta L θ←θ−η∇θL

其中：

params      模型参数
lr          learning rate
param.grad  参数梯度

---

8.1 torch.no_grad()

with torch.no_grad():

意思：

这部分操作不要记录到自动求导计算图中。

更新参数本身不需要再被求导。

---

8.2 zero_grad()

PyTorch 中梯度默认会累加。

因此每次新的 batch 训练之前通常需要：

optimizer.zero_grad()

清空旧梯度。

---

9. PyTorch 优化器

trainer = torch.optim.SGD(
    net.parameters(),
    lr=0.1
)

这里：

net.parameters()

自动获取模型中所有需要训练的参数。

标准训练流程：

optimizer.zero_grad()

y_hat = net(X)

loss = loss_fn(y_hat, y)

loss.backward()

optimizer.step()

对应：

清梯度
↓
forward
↓
计算 loss
↓
backward
↓
optimizer.step()
↓
更新参数

---

10. net.train() 与 net.eval()

net.train()

进入训练模式。

net.eval()

进入评估模式。

这对以下模块尤其重要：

Dropout
BatchNorm

例如 Dropout：

train 模式
→ 随机丢弃部分神经元

eval 模式
→ 不再随机丢弃

---

11. Dropout

手写 Dropout：

mask = (torch.rand(X.shape) > dropout).float()

return mask * X / (1.0 - dropout)

其中：

torch.rand(X.shape)

给每个元素生成一个 (0\sim1) 的随机数。

然后：

> dropout

生成：

True / False

再通过：

.float()

变成：

1 / 0

因此：

mask = 哪些神经元被保留

---

11.1 assert

assert 0 <= dropout <= 1

表示：

如果 dropout 不在 (0,1) 之间，程序立即报错。

---

12. argmax

y_hat.argmax(axis=1)

表示：

对每一行，找最大元素所在的位置。

例如：

[0.1, 0.2, 0.7]

执行 argmax：

2

表示预测类别：

class 2

---

13. Softmax 和 CrossEntropyLoss

Softmax：

y^j=eoj∑keok \hat y_j= \frac{e^{o_j}} {\sum_k e^{o_k}} y^j=∑keokeoj

作用：

把 logits 转换成总和为 1 的概率分布。

---

13.1 nn.CrossEntropyLoss

PyTorch：

nn.CrossEntropyLoss()

一般直接输入：

logits

因为它内部已经包含：

log-softmax
+
negative log likelihood

所以一般不需要：

softmax(logits)

之后再送进 CrossEntropyLoss。

---

13.2 reduction='none'

nn.CrossEntropyLoss(reduction='none')

表示：

保留 batch 中每个样本自己的 loss，不自动求平均或求和。

---

14. torchvision、FashionMNIST 与 ToTensor

from torchvision import datasets, transforms

torchvision 是 PyTorch 针对计算机视觉提供的库，包括：

图像数据集
图像预处理
常用视觉模型

---

14.1 ToTensor

trans = transforms.ToTensor()

作用：

将 PIL Image 或 ndarray 转换成 PyTorch Tensor。

图像通常变成：

[C, H, W]

格式。

---

14.2 FashionMNIST

datasets.FashionMNIST(
    root="../data",
    train=True,
    transform=trans,
    download=True
)

参数含义：

root       数据存在哪里
train      True=训练集，False=测试集
transform  对图像做什么预处理
download   如果没有数据是否自动下载

---

14.3 mnist_train[0][0]

mnist_train[0]

返回：

(image, label)

因此：

mnist_train[0][0]

是第 0 个样本的图片。

而：

mnist_train[0][1]

是第 0 个样本的 label。

---

15. DataLoader

train_iter = data.DataLoader(
    mnist_train,
    batch_size=256,
    shuffle=True,
    num_workers=4
)

参数：

dataset      数据集
batch_size   一次读取多少样本
shuffle      每个 epoch 是否打乱
num_workers  使用多少进程读取数据

DataLoader 是一个 iterable。

因此：

for X, y in train_iter:

每次获得一个 batch。

---

16. d2l.load_data_fashion_mnist

train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

这不是 PyTorch 原生函数。

它是 D2L 自己封装的辅助函数。

内部本质仍然是：

FashionMNIST
+
Transform
+
DataLoader

最后直接返回：

train DataLoader
test DataLoader

---

17. from d2l import torch as d2l

from d2l import torch as d2l

不是：

把 PyTorch 的 torch 改名为 d2l。

而是：

从 d2l 包中导入针对 PyTorch 实现的工具模块，并命名为 d2l。

所以：

d2l.sgd
d2l.Animator
d2l.load_data_fashion_mnist

都是 D2L 自己提供的辅助函数。

---

18. 常见参数初始化方法

均匀分布

torch.rand(...)

生成：

U(0,1) U(0,1) U(0,1)

标准正态

torch.randn(...)

生成：

N(0,1) N(0,1) N(0,1)

指定正态分布

torch.normal(mean, std, size)

生成：

N(μ,σ2) N(\mu,\sigma^2) N(μ,σ2)

PyTorch 初始化工具

nn.init.normal_(W, mean=0, std=0.01)

表示：

用正态分布原地初始化 W。

---

19. Xavier 初始化

Xavier 初始化的核心思想：

尽量让网络每层输入和输出的方差保持稳定。

典型形式：

σ2=2nin+nout\sigma^2 =\frac{2}{n_{in}+n_{out}} σ2=nin+nout2

目的之一：

缓解梯度爆炸
缓解梯度消失

---

20. Weight Decay / L2 正则化

损失函数：

L+λ∥w∥22 L+\lambda\|w\|_2^2 L+λ∥w∥22

代码可能写成：

loss(net(X), y) + lambd * l2_penalty(w)

作用：

不允许权重无限增大，从而限制模型复杂度。

PyTorch 中：

torch.optim.SGD(
    params,
    lr=...,
    weight_decay=wd
)

可以直接加入 weight decay。

---

21. CPU、GPU、CUDA

CPU：

通用计算
少量较强核心

GPU：

大量并行计算核心
特别适合矩阵运算

CUDA：

NVIDIA GPU 使用的并行计算平台。

判断 CUDA 是否可用：

torch.cuda.is_available()

---

21.1 .to(device)

device = torch.device("cuda")

net = net.to(device)
X = X.to(device)
y = y.to(device)

模型和数据必须位于同一个 device。

---

22. Batch Size 为什么影响显存？

batch size：

B B B

不会改变模型参数数量。

但会改变：

输入数据
activation
中间张量
反向传播缓存

的大小。

因此一般：

B↑⇒显存占用↑ B\uparrow \Rightarrow 显存占用\uparrow B↑⇒显存占用↑

---

23. nn.Conv2d

标准形式：

nn.Conv2d(
    Cin,
    Cout,
    kernel_size,
    stride,
    padding
)

例如：

nn.Conv2d(
    3,
    64,
    kernel_size=3,
    padding=1,
    stride=2
)

翻译：

输入通道 = 3
输出通道 = 64
kernel = 3×3
每次移动 = 2
边缘每边补 = 1

---

24. Conv2d 的 weight shape

conv.weight.shape

为：

Cout,Cin,Kh,Kw\] \\boxed{ \[C_{out},C_{in},K_h,K_w\] } \[Cout,Cin,Kh,Kw

例如：

nn.Conv2d(3, 64, 3)

weight shape：

[64, 3, 3, 3]

---

24.1 Conv2d 参数数量

一个输出通道需要：

CinKhKw C_{in}K_hK_w CinKhKw

个权重。

一共有：

Cout C_{out} Cout

个输出通道。

因此：

CoutCinKhKw C_{out}C_{in}K_hK_w CoutCinKhKw

个 weight。

再加：

Cout C_{out} Cout

个 bias。

最终：

CoutCinKhKw+Cout \boxed{ C_{out}C_{in}K_hK_w+C_{out} } CoutCinKhKw+Cout

---

25. CNN 和 MLP 参数的统一理解

MLP：

Din→Dout D_{in}\rightarrow D_{out} Din→Dout

参数量：

DoutDin+Dout D_{out}D_{in}+D_{out} DoutDin+Dout

CNN：

一个输出通道需要读取：

CinKhKw C_{in}K_hK_w CinKhKw

个局部输入特征。

所以：

CoutCinKhKw+Cout \boxed{ C_{out}C_{in}K_hK_w+C_{out} } CoutCinKhKw+Cout

可以理解成：

CNN 把局部 patch 展平后，在每个位置执行一套共享的线性变换。

---

26. Padding

padding=1

表示：

图片的每一边补 1。

经典组合：

kernel_size=3
padding=1
stride=1

可以保持空间尺寸不变。

---

27. Stride

stride=2

表示：

卷积核每次移动 2 格。

作用：

降低空间分辨率
完成下采样
减少计算量
扩大后续感受野

---

28. 多输入通道卷积：Cin → 1

D2L 手写代码：

def corr2d_multi_in(X, K):
    return sum(
        d2l.corr2d(x, k)
        for x, k in zip(X, K)
    )

这里：

X.shape = [Cin, H, W]
K.shape = [Cin, Kh, Kw]

zip(X, K) 会依次配对：

X[0] ↔ K[0]
X[1] ↔ K[1]
...

每个输入通道分别做卷积，然后：

sum(...)

把所有结果相加。

最终：

Cin→1 \boxed{ C_{in}\rightarrow1 } Cin→1

个输出通道。

---

29. torch.stack

假设：

A.shape = [2, 2]
B.shape = [2, 2]
C.shape = [2, 2]

执行：

torch.stack([A, B, C], dim=0)

结果：

[3, 2, 2]

作用：

新增一个维度，把若干相同 shape 的 Tensor 摞起来。

---

30. 多输入、多输出通道：Cin → Cout

完整 kernel：

K.shape=Cout,Cin,Kh,Kw \boxed{ K.shape=C_{out},C_{in},K_h,K_w } K.shape=Cout,Cin,Kh,Kw

D2L：

def corr2d_multi_in_out(X, K):
    return torch.stack(
        [corr2d_multi_in(X, k) for k in K],
        0
    )

逻辑：

每一套 [Cin, Kh, Kw] kernel
↓
生成一个 output channel

重复 Cout 次
↓
得到 Cout 个输出通道

stack
↓
[Cout, Hout, Wout]

---

31. 1×1 Conv

nn.Conv2d(
    Cin,
    Cout,
    kernel_size=1
)

因为 kernel 只有：

1×1 1\times1 1×1

所以完全不看空间邻居。

它只在每一个空间位置执行：Cin→Cout C_{in}\rightarrow C_{out} Cin→Cout

的线性变换。

因此：

1×1 Conv=在每个空间位置共享执行的Linear \boxed{1\times1\ Conv= 在每个空间位置共享执行的 Linear } 1×1 Conv=在每个空间位置共享执行的Linear

例如：

[B, 64, H, W]
↓ Conv2d(64, 16, 1)
[B, 16, H, W]

常见作用：

channel mixing
通道升维
通道降维
bottleneck

这也是 FNO 中经常出现 1×1 conv 的原因之一。

---

32. Attention 数据生成

例如：

n_train = 50

x_train, _ = torch.sort(
    torch.rand(n_train) * 5
)

表示：

在 0∼50\sim50∼5 之间随机生成 50 个训练输入，并排序。

定义真实函数：

def f(x):
    return 2 * torch.sin(x) + x**0.8

加入噪声：

y_train = (
    f(x_train)
    + torch.normal(0, 0.5, (n_train,))
)

数学上就是：

Y=f(X)+ε Y=f(X)+\varepsilon Y=f(X)+ε
其中：
ε∼N(0,0.52) \varepsilon\sim N(0,0.5^2) ε∼N(0,0.52)

33. **kwargs

def __init__(self, **kwargs):

表示：

把额外的关键字参数统一打包成 dictionary。

例如：

foo(a=1, b=2)

那么：

kwargs

类似：

{
    "a": 1,
    "b": 2
}

然后：

super().__init__(**kwargs)

把这些参数继续传给父类。

---

34. nn.Parameter

self.w = nn.Parameter(
    torch.rand(1)
)

表示：

将这个 Tensor 注册成模型需要训练的参数。

因此：

net.parameters()

可以自动找到它。

优化器也会更新它。

---

35. repeat_interleave

例如：

queries = torch.tensor([1., 4.])

执行：

queries.repeat_interleave(3)

得到：

[1, 1, 1, 4, 4, 4]

也就是：

每个元素连续重复 3 次。

Attention 中是为了让：

每一个 query

都可以与：

多个 key

逐个比较。

例如：

queries
[Q]

经过：

repeat_interleave(K)
.reshape(Q, K)

得到：

[Q, K]

---

36. Softmax 的 dim

假设：

scores.shape = [Q, K]

执行：

softmax(scores, dim=1)

表示：

每一个 query 都对自己的 K 个 keys 单独进行 softmax。

因此每一行满足：

∑kαqk=1 \sum_k\alpha_{qk}=1 k∑αqk=1

这里：

dim=1

对应：

key 这一维。

---

37. torch.bmm

bmm：

Batch Matrix Multiplication，批量矩阵乘法。

例如：

A.shape = [Q, 1, K]
B.shape = [Q, K, 1]

执行：

torch.bmm(A, B)

相当于分别执行 (Q) 次：

1,K\]×\[K,1\] \[1,K\]\\times\[K,1\] \[1,K\]×\[K,1

最终：

[Q, 1, 1]

Attention 中它实际上在计算：

∑iαivi \sum_i\alpha_i v_i i∑αivi

38. Nadaraya-Watson Attention 的完整 Shape

初始：

queries
[Q]

keys
[Q, K]

values
[Q, K]

query 扩展：

queries
[Q]
↓ repeat_interleave + reshape
[Q, K]

与 keys 做差：

queries - keys
→ [Q, K]

计算 attention score：

[Q, K]

然后：

softmax(..., dim=1)

得到：

attention_weights
[Q, K]

接着：

attention_weights.unsqueeze(1)

得到：

[Q, 1, K]

而：

values.unsqueeze(-1)

得到：

[Q, K, 1]

最后：

torch.bmm(...)

得到：

[Q, 1, 1]

再：

.reshape(-1)

变成：

[Q]

---

39. Attention 代码的数学逻辑

Nadaraya-Watson Attention 的 score：

sqi=−12((q−ki)w)2s_{qi}= -\frac12((q-k_i)w)^2 sqi=−21((q−ki)w)2

Softmax：

αqi=esqi∑jesqj \alpha_{qi}= \frac{e^{s_{qi}}} {\sum_j e^{s_{qj}}} αqi=∑jesqjesqi

最终：

y^q=∑iαqivi \hat y_q= \sum_i\alpha_{qi}v_i y^q=i∑αqivi

所以完整逻辑：

Query
↓
和所有 Key 比较
↓
得到 similarity score
↓
Softmax
↓
Attention Weights
↓
加权 Values
↓
Output

---

40. 我现在读深度学习代码最有用的习惯

以前我总想记：

这个函数有几个参数？

第二个参数是什么？

这个 API 到底怎么背？

后来发现更有效的方法是：

每看到一行代码，先问：

① 作者现在想干什么？

例如：

torch.normal(...)

作者可能是在：

生成人工噪声。

---

② 输入 Tensor 的 Shape 是什么？

例如：

X.shape = [B, Cin, H, W]

---

③ 执行完之后 Shape 变成什么？

例如：

nn.Conv2d(32, 64, 3)

脑子中应该直接形成：

[B, 32, H, W]
↓
读取 32 个输入 feature maps
↓
每次观察 3×3 局部区域
↓
生成 64 个新的 feature maps
↓
[B, 64, H', W']

再比如 Attention：

Query
↓
与所有 Keys 比较
↓
scores
↓
softmax
↓
attention weights
↓
对 Values 加权
↓
output

这样才能逐渐从：

"认识 PyTorch API"

过渡到：

"读懂模型真正的数据流和计算逻辑"。

---

总结：目前形成的代码阅读框架

现在看到一段神经网络代码，可以优先从三个层次阅读：

第一层：网络结构
有哪些 layer / block / module？

第二层：Tensor Shape
每一步：
输入 shape → 输出 shape

第三层：训练逻辑
forward
→ loss
→ backward
→ optimizer.step()

对于复杂模型，也尽量先不要陷进某一个 API。

先找：

输入是什么
↓
shape 是什么
↓
经过哪个模块
↓
shape 如何变化
↓
模块承担什么数学功能
↓
最终输出是什么

当这些关系能够在脑海中形成一张数据流图以后，代码本身就会从"一堆 Python"逐渐变成"数学模型的实现"。