6. 各类库梳理

PyTorch整体架构系统教程¶

目标读者

• 已经熟悉 Python。
• 已经会写基础 torch.Tensor 运算、矩阵乘法、embedding 计算、简单 loss。
• 但对 PyTorch 的框架结构、模块边界、训练流程协作关系还不清楚。

学习目标

• 从全局视角理解 PyTorch 的整体设计理念。
• 搞清楚 torch、torch.nn、torch.nn.functional、torch.optim、torch.autograd、torch.utils 等模块分别负责什么。
• 建立“数据如何进入模型、梯度如何反传、参数如何更新”的完整心智模型。
• 从“会写局部算法代码”提升到“理解框架如何组织这些代码”。

Notebook 导航¶

1. PyTorch 的整体架构概览
2. Tensor：PyTorch 的核心数据结构
3. 自动求导系统：torch.autograd
4. 神经网络模块系统：torch.nn
5. torch.nn.functional 的作用
6. 损失函数系统
7. 优化器系统：torch.optim
8. 数据系统：Dataset 与 DataLoader
9. GPU 与设备管理
10. PyTorch 的完整训练流程
11. 常见 API 分类总结
12. 从“写算法”到“理解框架”的思维转变

In [1]:

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, TensorDataset

print("torch version:", torch.__version__)
print("cuda available:", torch.cuda.is_available())

from __future__ import annotations import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset, TensorDataset print("torch version:", torch.__version__) print("cuda available:", torch.cuda.is_available())

torch version: 2.4.1+cu124
cuda available: True

1. PyTorch 的整体架构概览¶

PyTorch 可以理解成一个围绕 Tensor + 自动求导 + 神经网络组件 + 训练工具链 组织起来的深度学习框架。

1.1 各核心模块分别解决什么问题¶

模块	主要职责	你在训练中会怎么用
torch	基础张量库，提供 Tensor、数学运算、随机数、设备管理等	写张量运算、矩阵乘法、reshape、索引
torch.nn	神经网络模块化封装，提供 nn.Module、常见层、常见 loss	定义模型结构、注册参数
torch.nn.functional	无状态的函数式操作	写 relu、softmax、cross_entropy
torch.optim	优化器系统	用 SGD/Adam/AdamW 更新模型参数
torch.autograd	自动求导引擎	让 backward() 自动计算梯度
torch.utils	通用工具集合	重点是 torch.utils.data 中的 Dataset/DataLoader
torch.cuda	CUDA 设备与 GPU 相关能力	把 Tensor/模型搬到 GPU，检查显存和设备
torch.distributed	分布式训练能力	多卡、多机训练
torch.backends	后端行为与性能选项	控制 cudnn / mps / mkldnn 等后端行为

1.2 PyTorch 的设计理念¶

1. 以 Tensor 为中心：模型参数、输入数据、梯度，本质上都是 Tensor。
2. 以动态图为核心：前向计算时动态构图，因此写法更接近普通 Python 程序。
3. 以模块化组合为主要抽象：层、模型、loss、优化器、数据管道被拆成相互协作的组件。
4. Python 前端 + 高性能后端：用户在 Python 中写逻辑，底层调用高性能 C/CUDA 后端执行。

1.3 PyTorch 训练流程在框架中的位置关系图¶

原始样本
  -> Dataset                         (torch.utils.data.Dataset)
  -> DataLoader                      (torch.utils.data.DataLoader)
  -> batch Tensor                    (torch.Tensor)
  -> Model.forward()                 (torch.nn.Module)
  -> prediction / logits             (torch.Tensor)
  -> Loss                            (torch.nn 或 torch.nn.functional)
  -> backward()                      (torch.autograd)
  -> parameter.grad                  (Tensor 上的梯度)
  -> optimizer.step()                (torch.optim)
  -> 参数更新后的 Model

这张图揭示了 PyTorch 的核心思想：

• torch 负责“数据和运算”。
• torch.nn 负责“模型和可学习参数”。
• torch.autograd 负责“梯度是怎么来的”。
• torch.optim 负责“参数如何更新”。
• torch.utils.data 负责“数据如何按 batch 喂给模型”。

In [2]:

module_roles = {
    "torch": "Tensor 与基础运算",
    "torch.nn": "模型层与参数容器",
    "torch.nn.functional": "无状态函数式算子",
    "torch.autograd": "自动求导",
    "torch.optim": "参数更新",
    "torch.utils.data": "数据管道",
    "torch.cuda": "GPU 设备管理",
}

for name, role in module_roles.items():
    print(f"{name:20s} -> {role}")

module_roles = { "torch": "Tensor 与基础运算", "torch.nn": "模型层与参数容器", "torch.nn.functional": "无状态函数式算子", "torch.autograd": "自动求导", "torch.optim": "参数更新", "torch.utils.data": "数据管道", "torch.cuda": "GPU 设备管理", } for name, role in module_roles.items(): print(f"{name:20s} -> {role}")

torch                -> Tensor 与基础运算
torch.nn             -> 模型层与参数容器
torch.nn.functional  -> 无状态函数式算子
torch.autograd       -> 自动求导
torch.optim          -> 参数更新
torch.utils.data     -> 数据管道
torch.cuda           -> GPU 设备管理

2. Tensor：PyTorch 的核心数据结构¶

torch.Tensor 是 PyTorch 的核心。你看到的几乎一切训练对象，最终都能落到 Tensor 上：

• 输入数据是 Tensor。
• 模型参数是 Tensor。
• loss 是一个 Tensor。
• 梯度也是 Tensor。

2.1 Tensor 与 NumPy 的区别¶

它们都支持多维数组和广播，但 PyTorch 的 Tensor 更强调：

• 自动求导：Tensor 可以记录计算历史。
• GPU 支持：Tensor 可以在 CUDA 设备上运行。
• 与神经网络训练系统天然集成：参数、优化器、模块系统都围绕 Tensor 设计。

2.2 需要重点理解的几个属性¶

• shape：张量形状。
• dtype：数据类型，如 float32、int64。
• device：张量所在设备，如 CPU / CUDA。
• requires_grad：是否需要自动求梯度。

In [3]:

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
y = torch.tensor([10.0, 20.0])

print("x =")
print(x)
print("shape:", x.shape)
print("dtype:", x.dtype)
print("device:", x.device)

print()
print("Broadcasting result:")
print(x + y)

np_view = x.numpy()
print()
print("Converted to numpy type:", type(np_view))

grad_tensor = torch.randn(2, 3, requires_grad=True)
print()
print("requires_grad:", grad_tensor.requires_grad)

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32) y = torch.tensor([10.0, 20.0]) print("x =") print(x) print("shape:", x.shape) print("dtype:", x.dtype) print("device:", x.device) print() print("Broadcasting result:") print(x + y) np_view = x.numpy() print() print("Converted to numpy type:", type(np_view)) grad_tensor = torch.randn(2, 3, requires_grad=True) print() print("requires_grad:", grad_tensor.requires_grad)

x =
tensor([[1., 2.],
        [3., 4.]])
shape: torch.Size([2, 2])
dtype: torch.float32
device: cpu

Broadcasting result:
tensor([[11., 22.],
        [13., 24.]])

Converted to numpy type: <class 'numpy.ndarray'>

requires_grad: True

2.3 Broadcasting 的意义¶

广播不是深度学习专属概念，但在 PyTorch 中极其常见。它让你可以不手动复制数据，就完成按维度对齐的运算。

例如：

• [batch, hidden] + [hidden]
• [batch, seq, hidden] * [hidden]

这也是很多 loss、归一化、embedding 后处理写起来简洁的原因。

3. 自动求导系统（torch.autograd）¶

很多人知道 loss.backward() 可以求梯度，但不知道它为什么能工作。

3.1 关键机制：计算图¶

当你对 requires_grad=True 的 Tensor 做运算时，PyTorch 会在后台记录：

• 这个结果 Tensor 是由哪些 Tensor 计算得到的。
• 每一步用了什么运算。

这些依赖关系构成一张 计算图（computation graph）。

当前向结束后，如果你对最终标量调用 backward()，PyTorch 就会沿着这张图反向传播，自动用链式法则把每个参数的梯度算出来。

3.2 你需要掌握的几个接口¶

• requires_grad=True：告诉 PyTorch 这个 Tensor 需要梯度。
• loss.backward()：从当前标量向前回传，计算图中相关 Tensor 的梯度。
• tensor.grad：查看叶子 Tensor 的梯度。
• detach()：返回一个“断开计算图”的 Tensor，不再继续追踪梯度。

In [4]:

w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
x = torch.tensor(3.0)

y = w * x + b
loss = (y - 10) ** 2

loss.backward()

print("loss:", loss.item())
print("dloss/dw:", w.grad.item())
print("dloss/db:", b.grad.item())

detached_y = y.detach()
print("detached_y requires_grad:", detached_y.requires_grad)

w = torch.tensor(2.0, requires_grad=True) b = torch.tensor(1.0, requires_grad=True) x = torch.tensor(3.0) y = w * x + b loss = (y - 10) ** 2 loss.backward() print("loss:", loss.item()) print("dloss/dw:", w.grad.item()) print("dloss/db:", b.grad.item()) detached_y = y.detach() print("detached_y requires_grad:", detached_y.requires_grad)

loss: 9.0
dloss/dw: -18.0
dloss/db: -6.0
detached_y requires_grad: False

3.3 为什么 PyTorch 可以自动求梯度¶

因为 PyTorch 不只是“做数值计算”，还在做两件事：

1. 记录前向依赖关系：结果 Tensor 知道自己从哪里来。
2. 为每个算子实现反向规则：如乘法、加法、矩阵乘法、卷积等，都有对应的梯度传播逻辑。

所以你写的是前向代码，但框架在背后维护了“如何反向”的信息。

这就是为什么研究者只写 forward，仍然能训练复杂模型。

4. 神经网络模块系统（torch.nn）¶

4.1 nn.Module 的设计思想¶

nn.Module 是 PyTorch 中“模型”最核心的抽象。它解决的不是数值运算，而是 如何组织可学习参数和子模块。

一个 nn.Module 通常做三件事：

• 保存可学习参数。
• 保存子模块，例如若干层 Linear、Conv2d、LayerNorm。
• 定义前向逻辑 forward()。

4.2 为什么模型通常必须继承 nn.Module¶

因为框架需要统一管理这些能力：

• 自动收集参数：model.parameters()
• 切换训练/推理模式：model.train() / model.eval()
• 设备迁移：model.to(device)
• 保存和加载：state_dict() / load_state_dict()

如果不继承 nn.Module，这些标准化能力就无法自动获得。

4.3 parameters() 如何工作¶

当你把 nn.Linear、nn.Embedding 等子模块赋值给 self.xxx 时，nn.Module 会自动注册它们内部的参数。

所以：

• self.linear = nn.Linear(...) 会被框架识别为子模块。
• 这个子模块里的 weight、bias 会自动进入 model.parameters()。

4.4 forward() 的作用¶

forward() 只负责定义“输入如何变成输出”。

训练时你通常写：

logits = model(x)

表面上你在调用 model(x)，本质上框架会转到 forward()，同时保留模块钩子、模式切换等机制。

4.5 常见层为什么属于 torch.nn¶

因为这些层不仅是一个运算，更是一个“带参数或带模块语义的组件”：

• nn.Linear：线性层，内部有 weight / bias
• nn.Conv2d：卷积层，内部有卷积核参数
• nn.Embedding：离散 id 到向量的查表层，内部有 embedding matrix
• nn.LayerNorm：归一化层，常常带可学习缩放/偏置
• nn.Dropout：训练/推理行为不同，所以需要模块语义
• nn.ReLU：虽然无参数，但作为层组合时可读性高，也可放进 nn.Sequential

In [5]:

class TinyMLP(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.norm = nn.LayerNorm(hidden_dim)
        self.dropout = nn.Dropout(p=0.1)
        self.fc2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = F.relu(x)
        x = self.norm(x)
        x = self.dropout(x)
        return self.fc2(x)

model = TinyMLP(in_dim=8, hidden_dim=16, out_dim=4)
sample = torch.randn(2, 8)
out = model(sample)

print(model)
print("output shape:", out.shape)
print()
print("Named parameters:")
for name, param in model.named_parameters():
    print(f"{name:20s}", tuple(param.shape))

class TinyMLP(nn.Module): def __init__(self, in_dim: int, hidden_dim: int, out_dim: int): super().__init__() self.fc1 = nn.Linear(in_dim, hidden_dim) self.norm = nn.LayerNorm(hidden_dim) self.dropout = nn.Dropout(p=0.1) self.fc2 = nn.Linear(hidden_dim, out_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.fc1(x) x = F.relu(x) x = self.norm(x) x = self.dropout(x) return self.fc2(x) model = TinyMLP(in_dim=8, hidden_dim=16, out_dim=4) sample = torch.randn(2, 8) out = model(sample) print(model) print("output shape:", out.shape) print() print("Named parameters:") for name, param in model.named_parameters(): print(f"{name:20s}", tuple(param.shape))

TinyMLP(
  (fc1): Linear(in_features=8, out_features=16, bias=True)
  (norm): LayerNorm((16,), eps=1e-05, elementwise_affine=True)
  (dropout): Dropout(p=0.1, inplace=False)
  (fc2): Linear(in_features=16, out_features=4, bias=True)
)
output shape: torch.Size([2, 4])

Named parameters:
fc1.weight           (16, 8)
fc1.bias             (16,)
norm.weight          (16,)
norm.bias            (16,)
fc2.weight           (4, 16)
fc2.bias             (4,)

5. torch.nn.functional 的作用¶

torch.nn.functional 常被记作 F。它提供的是 函数式（functional） 的操作。

典型疑问是：为什么既有 nn.ReLU()，又有 F.relu()？

5.1 Module 版本 vs Functional 版本¶

对比项	Module 版本	Functional 版本
形式	nn.ReLU()	F.relu(x)
是否是对象	是，一个模块实例	否，一个普通函数
是否便于注册进模型结构	是	否
是否有参数状态	有些模块有，有些没有	一般无状态
适合场景	模型结构清晰、放进 nn.Sequential	临时调用、灵活写 forward

直观理解：

• nn.Module 版本强调“这是模型结构的一部分”。
• F.xxx 版本强调“我现在只想做一次运算”。

因此很多研究代码里会混用：

• 线性层、卷积层、embedding 常写成 nn.Linear / nn.Conv2d。
• 激活函数、loss、归一化操作常直接写 F.relu / F.cross_entropy / F.normalize。

In [6]:

x = torch.tensor([[-1.0, 0.5, 2.0]])

relu_module = nn.ReLU()
print("Module ReLU:", relu_module(x))
print("Functional ReLU:", F.relu(x))

logits = torch.tensor([[2.0, 1.0, 0.1]])
probs = F.softmax(logits, dim=-1)
print()
print("softmax:", probs)

features = torch.tensor([[3.0, 4.0]])
print("normalize:", F.normalize(features, p=2, dim=-1))

target = torch.tensor([0])
print("cross_entropy:", F.cross_entropy(logits, target).item())

x = torch.tensor([[-1.0, 0.5, 2.0]]) relu_module = nn.ReLU() print("Module ReLU:", relu_module(x)) print("Functional ReLU:", F.relu(x)) logits = torch.tensor([[2.0, 1.0, 0.1]]) probs = F.softmax(logits, dim=-1) print() print("softmax:", probs) features = torch.tensor([[3.0, 4.0]]) print("normalize:", F.normalize(features, p=2, dim=-1)) target = torch.tensor([0]) print("cross_entropy:", F.cross_entropy(logits, target).item())

Module ReLU: tensor([[0.0000, 0.5000, 2.0000]])
Functional ReLU: tensor([[0.0000, 0.5000, 2.0000]])

softmax: tensor([[0.6590, 0.2424, 0.0986]])
normalize: tensor([[0.6000, 0.8000]])
cross_entropy: 0.4170299470424652

一、底层原理区别

torch.nn.functional 本质是 算子函数库。 每个函数只是执行一次前向计算，本身不保存任何参数或运行状态。如果计算需要权重或偏置，必须作为参数显式传入。例如：

F.linear(x, weight, bias)

函数只负责计算：

[ y = xW^T + b ]

但 weight、bias 并不属于这个函数。

torch.nn 提供的是 Module（网络层对象）。 Module 内部会 保存参数并注册到模型中，这些参数会自动参与：

• 反向传播
• optimizer 更新
• 模型保存 / 加载

例如：

layer = nn.Linear(10, 5)
y = layer(x)

这里 weight 和 bias 是 layer 的内部成员：

layer.weight
layer.bias

因此：

functional = 纯计算函数 nn.Module = 带参数管理的网络层封装

二、两者关系

大多数 nn.Module 在内部其实调用的就是 functional。

例如逻辑结构：

nn.Linear.forward()
        ↓
F.linear(x, weight, bias)

也就是说：

functional 负责计算
Module 负责参数管理

三、实际使用习惯

工程实践通常遵循一个简单原则：

有参数的层 → 用 nn.Module

例如：

• nn.Linear
• nn.Conv2d
• nn.BatchNorm
• nn.Embedding

因为需要自动管理参数。

无参数操作 → 用 functional

例如：

• F.relu
• F.softmax
• F.dropout
• F.cross_entropy

这些只是计算，不需要保存权重。

6. 损失函数系统¶

损失函数在 PyTorch 中通常有两种入口：

• torch.nn 里的模块形式，如 nn.CrossEntropyLoss()
• torch.nn.functional 里的函数形式，如 F.cross_entropy(...)

6.1 常见损失¶

• nn.CrossEntropyLoss：多分类任务常用。
• nn.BCELoss：输入通常是概率，数值稳定性不如 logits 版本。
• nn.BCEWithLogitsLoss：把 sigmoid 和 BCE 合在一起，数值更稳定。
• nn.MSELoss：回归任务常见。

6.2 为什么很多人更喜欢 F.cross_entropy¶

因为它：

• 写起来更短，尤其在研究代码中很常见。
• 不需要提前实例化对象。
• 更适合在 forward 外部灵活控制计算逻辑。

但 nn.CrossEntropyLoss() 也有优势：

• 当你想把 loss 配置封装成模块对象时更统一。
• 某些场景下便于和配置系统、训练器框架集成。

核心不是“谁更高级”，而是：

• loss 是否需要作为对象长期保存。
• 你更偏向函数式研究写法，还是模块式工程写法。

In [7]:

logits = torch.tensor([[2.5, 0.3], [0.1, 1.2]])
class_target = torch.tensor([0, 1])

ce_module = nn.CrossEntropyLoss()
ce_loss_1 = ce_module(logits, class_target)
ce_loss_2 = F.cross_entropy(logits, class_target)

probs = torch.tensor([[0.8], [0.2]])
binary_target = torch.tensor([[1.0], [0.0]])
bce_loss = nn.BCELoss()(probs, binary_target)

raw_logits = torch.tensor([[1.5], [-1.2]])
bce_logits_loss = nn.BCEWithLogitsLoss()(raw_logits, binary_target)

pred = torch.tensor([[2.0], [3.5]])
truth = torch.tensor([[1.5], [4.0]])
mse_loss = nn.MSELoss()(pred, truth)

print("CrossEntropyLoss(module):", ce_loss_1.item())
print("cross_entropy(functional):", ce_loss_2.item())
print("BCELoss:", bce_loss.item())
print("BCEWithLogitsLoss:", bce_logits_loss.item())
print("MSELoss:", mse_loss.item())

logits = torch.tensor([[2.5, 0.3], [0.1, 1.2]]) class_target = torch.tensor([0, 1]) ce_module = nn.CrossEntropyLoss() ce_loss_1 = ce_module(logits, class_target) ce_loss_2 = F.cross_entropy(logits, class_target) probs = torch.tensor([[0.8], [0.2]]) binary_target = torch.tensor([[1.0], [0.0]]) bce_loss = nn.BCELoss()(probs, binary_target) raw_logits = torch.tensor([[1.5], [-1.2]]) bce_logits_loss = nn.BCEWithLogitsLoss()(raw_logits, binary_target) pred = torch.tensor([[2.0], [3.5]]) truth = torch.tensor([[1.5], [4.0]]) mse_loss = nn.MSELoss()(pred, truth) print("CrossEntropyLoss(module):", ce_loss_1.item()) print("cross_entropy(functional):", ce_loss_2.item()) print("BCELoss:", bce_loss.item()) print("BCEWithLogitsLoss:", bce_logits_loss.item()) print("MSELoss:", mse_loss.item())

CrossEntropyLoss(module): 0.19620931148529053
cross_entropy(functional): 0.19620931148529053
BCELoss: 0.2231435477733612
BCEWithLogitsLoss: 0.23234784603118896
MSELoss: 0.25

7. 优化器系统（torch.optim）¶

自动求导只负责算出梯度，但 不会自动更新参数。参数更新由优化器负责。

7.1 optimizer 的作用¶

优化器接收一组参数，通常来自 model.parameters()，根据这些参数的梯度 param.grad 更新它们。

常见优化器：

• SGD：最基础的随机梯度下降，可配合 momentum。
• Adam：自适应学习率方法，研究中非常常见。
• AdamW：将权重衰减与 Adam 更合理地解耦，是 Transformer 等模型里常见默认选择。

7.2 训练中的标准三步¶

optimizer.zero_grad()
loss.backward()
optimizer.step()

解释如下：

• zero_grad()：把上一步残留的梯度清零。
• backward()：计算当前 batch 对参数的梯度。
• step()：按照优化算法更新参数。

为什么先清零？因为 PyTorch 默认是 梯度累加，不是自动覆盖。

In [8]:

linear = nn.Linear(3, 1)
optimizer = torch.optim.SGD(linear.parameters(), lr=0.1)

x = torch.randn(4, 3)
y = torch.randn(4, 1)

optimizer.zero_grad()
pred = linear(x)
loss = F.mse_loss(pred, y)
loss.backward()

print("grad norm before step:", linear.weight.grad.norm().item())
optimizer.step()
print("one optimizer step finished")

linear = nn.Linear(3, 1) optimizer = torch.optim.SGD(linear.parameters(), lr=0.1) x = torch.randn(4, 3) y = torch.randn(4, 1) optimizer.zero_grad() pred = linear(x) loss = F.mse_loss(pred, y) loss.backward() print("grad norm before step:", linear.weight.grad.norm().item()) optimizer.step() print("one optimizer step finished")

grad norm before step: 3.7122721672058105
one optimizer step finished

8. 数据系统（Dataset 与 DataLoader）¶

数据系统主要位于 torch.utils.data。

8.1 Dataset 是什么¶

Dataset 解决的是“单个样本如何被读取”的问题。你通常需要定义：

• __len__()：数据集大小。
• __getitem__(idx)：如何取第 idx 个样本。

8.2 DataLoader 是什么¶

DataLoader 解决的是“如何把 Dataset 组织成 batch 并高效迭代”的问题。它负责：

• batch 切分
• shuffle
• 多进程加载，num_workers
• 自动拼接 batch

8.3 为什么训练必须用 DataLoader¶

因为深度学习训练通常不是“一次喂完整个数据集”，而是：

• 按 batch 迭代
• 每个 batch 做一次前向、反向、更新
• 多个 epoch 重复

DataLoader 把这些机械工作标准化了。

In [9]:

class ToyDataset(Dataset):
    def __init__(self):
        self.features = torch.tensor([
            [1.0, 0.0],
            [0.0, 1.0],
            [1.0, 1.0],
            [0.0, 0.0],
        ])
        self.labels = torch.tensor([0, 1, 1, 0])

    def __len__(self) -> int:
        return len(self.features)

    def __getitem__(self, idx: int):
        return self.features[idx], self.labels[idx]

dataset = ToyDataset()
loader = DataLoader(dataset, batch_size=2, shuffle=True)

for batch_x, batch_y in loader:
    print("batch_x shape:", batch_x.shape, "batch_y:", batch_y)

class ToyDataset(Dataset): def __init__(self): self.features = torch.tensor([ [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.0, 0.0], ]) self.labels = torch.tensor([0, 1, 1, 0]) def __len__(self) -> int: return len(self.features) def __getitem__(self, idx: int): return self.features[idx], self.labels[idx] dataset = ToyDataset() loader = DataLoader(dataset, batch_size=2, shuffle=True) for batch_x, batch_y in loader: print("batch_x shape:", batch_x.shape, "batch_y:", batch_y)

batch_x shape: torch.Size([2, 2]) batch_y: tensor([1, 0])
batch_x shape: torch.Size([2, 2]) batch_y: tensor([0, 1])

9. GPU 与设备管理¶

在 PyTorch 里，Tensor 和模型都带有设备属性。

9.1 核心概念¶

• device = torch.device("cuda") 或 torch.device("cpu")
• tensor.to(device)：把 Tensor 移动到目标设备
• model.to(device)：把模型参数移动到目标设备

9.2 设备管理最重要的原则¶

参与同一次计算的 Tensor 和模型参数必须在同一设备上。

比如：

• 模型在 GPU，输入还在 CPU，就会报错。
• 标签忘了搬到 GPU，也常见报错。

所以训练代码里经常会看到：

x = x.to(device)
y = y.to(device)
model = model.to(device)

In [10]:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x = torch.randn(2, 3)
model = nn.Linear(3, 2)

x = x.to(device)
model = model.to(device)
out = model(x)

print("device:", device)
print("x device:", x.device)
print("model weight device:", model.weight.device)
print("out device:", out.device)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") x = torch.randn(2, 3) model = nn.Linear(3, 2) x = x.to(device) model = model.to(device) out = model(x) print("device:", device) print("x device:", x.device) print("model weight device:", model.weight.device) print("out device:", out.device)

device: cuda
x device: cuda:0
model weight device: cuda:0
out device: cuda:0

10. PyTorch 的完整训练流程¶

下面这段代码把前面的模块串起来，形成一个完整可读的训练闭环。

10.1 先看职责分工¶

• TensorDataset / DataLoader：来自 torch.utils.data
• Classifier / nn.Linear / nn.ReLU：来自 torch.nn
• F.cross_entropy：来自 torch.nn.functional
• loss.backward()：触发 torch.autograd
• torch.optim.Adam：来自 torch.optim
• tensor.to(device) / model.to(device)：设备管理属于 torch / torch.cuda

In [11]:

torch.manual_seed(7)

# 1) dataset: 构造一个可分的二分类玩具数据
features = torch.randn(256, 2)
labels = (features[:, 0] + 0.8 * features[:, 1] > 0).long()

dataset = TensorDataset(features, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 2) model: 继承 nn.Module 定义网络
class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 16),
            nn.ReLU(),
            nn.Linear(16, 2),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Classifier().to(device)

# 3) optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

# 4) training loop
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    total_loss = 0.0
    total_correct = 0
    total_samples = 0

    for batch_x, batch_y in dataloader:
        batch_x = batch_x.to(device)
        batch_y = batch_y.to(device)

        logits = model(batch_x)
        loss = F.cross_entropy(logits, batch_y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * batch_x.size(0)
        total_correct += (logits.argmax(dim=1) == batch_y).sum().item()
        total_samples += batch_x.size(0)

    avg_loss = total_loss / total_samples
    avg_acc = total_correct / total_samples
    print(f"epoch={epoch + 1} loss={avg_loss:.4f} acc={avg_acc:.4f}")

torch.manual_seed(7) # 1) dataset: 构造一个可分的二分类玩具数据 features = torch.randn(256, 2) labels = (features[:, 0] + 0.8 * features[:, 1] > 0).long() dataset = TensorDataset(features, labels) dataloader = DataLoader(dataset, batch_size=32, shuffle=True) # 2) model: 继承 nn.Module 定义网络 class Classifier(nn.Module): def __init__(self): super().__init__() self.net = nn.Sequential( nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Classifier().to(device) # 3) optimizer optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) # 4) training loop num_epochs = 5 for epoch in range(num_epochs): model.train() total_loss = 0.0 total_correct = 0 total_samples = 0 for batch_x, batch_y in dataloader: batch_x = batch_x.to(device) batch_y = batch_y.to(device) logits = model(batch_x) loss = F.cross_entropy(logits, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() * batch_x.size(0) total_correct += (logits.argmax(dim=1) == batch_y).sum().item() total_samples += batch_x.size(0) avg_loss = total_loss / total_samples avg_acc = total_correct / total_samples print(f"epoch={epoch + 1} loss={avg_loss:.4f} acc={avg_acc:.4f}")

epoch=1 loss=0.6657 acc=0.5586
epoch=2 loss=0.4811 acc=0.8789
epoch=3 loss=0.3534 acc=0.9570
epoch=4 loss=0.2605 acc=0.9648
epoch=5 loss=0.1977 acc=0.9883

10.2 用一句话概括训练循环¶

每一轮训练本质上都在重复以下闭环：

取一个 batch
  -> 模型前向得到预测
  -> 计算 loss
  -> backward 计算梯度
  -> optimizer.step 更新参数
  -> 进入下一个 batch

这就是 PyTorch 框架结构最核心的协作方式。

11. PyTorch 常见 API 分类总结¶

下表把最常见的 API 做一个结构化总结，帮助你把“常见写法”放回它所属的模块里。

分类	常见 API	所属模块	主要作用
Tensor 运算	torch.tensor, torch.randn, matmul, sum, view, reshape	torch	创建和变换张量
自动求导	requires_grad, backward, grad, detach	torch.autograd / Tensor 接口	记录计算图并计算梯度
网络层	nn.Linear, nn.Conv2d, nn.Embedding, nn.LayerNorm, nn.Dropout, nn.ReLU	torch.nn	定义可复用网络组件
函数式算子	F.relu, F.softmax, F.cross_entropy, F.normalize	torch.nn.functional	无状态前向运算
损失函数	nn.CrossEntropyLoss, nn.BCEWithLogitsLoss, nn.MSELoss	torch.nn	以模块形式组织 loss
优化器	torch.optim.SGD, torch.optim.Adam, torch.optim.AdamW	torch.optim	根据梯度更新参数
数据加载	Dataset, DataLoader, TensorDataset	torch.utils.data	管理样本、batch、shuffle
设备管理	torch.device, tensor.to(device), model.cuda()	torch / torch.cuda	CPU/GPU 切换
分布式	DistributedDataParallel 等	torch.distributed	多卡多机训练
后端控制	torch.backends.cudnn 等	torch.backends	控制后端性能与行为

12. 从“写算法”到“理解框架”的思维转变¶

很多人会写 loss、attention、embedding、矩阵乘法，但仍然不理解框架，原因通常不是数学不会，而是 没有建立组件分工意识。

12.1 常见误区¶

误区 1：把 PyTorch 理解成“很多 API 的集合”

实际上它更像一条协作流水线：

• 数据系统负责喂数据。
• 模型系统负责组织参数和前向逻辑。
• autograd 负责生成梯度。
• optimizer 负责更新参数。

误区 2：只关注单个算子，不关注对象之间如何配合

会写 cross_entropy 不等于理解训练。真正需要理解的是：

• 这个 loss 的输入从哪里来？
• 它为什么能反传到模型参数？
• 参数是谁收集的？
• 参数又是谁更新的？

误区 3：把模型理解成“一个函数”

在 PyTorch 中，模型不只是函数，更是：

• 参数容器
• 子模块树
• 训练/推理模式切换入口
• 状态保存载体

12.2 建立 PyTorch 的整体心智模型¶

你可以把 PyTorch 记成 5 个层次：

1. 数据层：Dataset / DataLoader
2. 表示层：Tensor
3. 模型层：nn.Module 与各种网络层
4. 求导层：autograd
5. 优化层：optim

最后把它串成一句话：

用 DataLoader 产生 batch Tensor，喂给 nn.Module 做 forward，
用 loss 衡量预测误差，autograd 负责 backward 求梯度，
optimizer 根据梯度更新参数。

当你能从这个角度看代码时，你就不再只是“会写一些 PyTorch 运算”，而是开始真正理解这个框架了。

In [12]:

framework_map = {
    "数据流": "Dataset -> DataLoader -> batch Tensor",
    "模型流": "nn.Module.forward -> logits",
    "训练流": "loss -> backward -> optimizer.step",
    "设备流": "cpu/cuda 之间迁移 Tensor 和模型",
}

for k, v in framework_map.items():
    print(f"{k}: {v}")

framework_map = { "数据流": "Dataset -> DataLoader -> batch Tensor", "模型流": "nn.Module.forward -> logits", "训练流": "loss -> backward -> optimizer.step", "设备流": "cpu/cuda 之间迁移 Tensor 和模型", } for k, v in framework_map.items(): print(f"{k}: {v}")

数据流: Dataset -> DataLoader -> batch Tensor
模型流: nn.Module.forward -> logits
训练流: loss -> backward -> optimizer.step
设备流: cpu/cuda 之间迁移 Tensor 和模型