【AI入门知识点】PyTorch 是什么？从 Tensor 到多 GPU 训练，小白也能真正入门
一、PyTorch 到底是什么（真正讲清楚）
小白版：
想象你要搭乐高。你可以自己烧塑料、开模具、做积木——这相当于用 C++ 从零写神经网络。你也可以买现成的乐高套装，里面有标准积木、说明书，你只需要按图纸拼——PyTorch 就是这套"深度学习乐高套装"。它把"矩阵运算"“自动算梯度”"模型组件"这些基础零件都给你准备好了，你只要像写普通 Python 代码一样把它们拼起来就行。你不用操心积木是怎么造出来的，只管搭你的城堡。
专业版：
PyTorch 是一个深度学习框架。但这句话等于没说，因为 TensorFlow、JAX、飞桨都是深度学习框架。
更准确的定义是：PyTorch 是一个由 Meta（Facebook）开源的、面向 Python 的张量计算库，它的核心设计原则是"让写深度学习代码就像写普通 Python 代码一样自然"。
拆开来看，PyTorch 做了三件事：
| 能力 | 说明 | 类比 |
|---|---|---|
| 张量计算 | 替代 NumPy，且能跑在 GPU 上 | NumPy 的 GPU 版本 |
| 自动求导 | 你只写前向计算，反向梯度自动算 | 计算器按等号 |
| 神经网络模块 | 提供现成的层、损失函数、优化器 | 搭积木的零件包 |
举个直白的例子：你用 NumPy 写了一个线性回归，想让它跑在 GPU 上、自动算梯度、用 Adam 优化——这需要你自己手写几百行反向传播代码。而用 PyTorch，核心代码就几行：
import torch
import torch.nn as nn
import torch.optim as optim
# 定义一个最简单的模型
model = nn.Linear(10, 1)
# 损失函数 + 优化器
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 一个训练步
x = torch.randn(32, 10) # 32 个样本，每个 10 维
y = torch.randn(32, 1) # 标签
pred = model(x) # 前向传播
loss = loss_fn(pred, y) # 计算损失
loss.backward() # 反向传播 —— 就这一行！
optimizer.step() # 更新参数
optimizer.zero_grad() # 清空梯度
这就是 PyTorch 的核心价值：把深度学习中重复的、底层的、易错的部分抽象掉，让你专注于模型设计本身。
二、为什么 AI 圈几乎都在用 PyTorch
小白版：
这就像当年智能手机大战：诺基亚（TensorFlow 1.x）功能强大但操作反人类，iPhone（PyTorch）一出来大家发现"原来写代码可以这么爽"。早期 TensorFlow 要求你先画一张"全局施工图"再施工，中途你想看看盖到哪了都不行；而 PyTorch 是你一边写代码一边就能看到结果，出错了马上能定位。这种"所见即所得"的体验让学术圈迅速倒戈，而工业界招聘时发现候选人都会 PyTorch，公司也就跟着转了。等 Google 回过神把 TensorFlow 改成类似 PyTorch 的风格（2.0），地盘已经丢了大半。
专业版：
根据 Papers with Code 的统计，2024 年 超过 80% 的顶会论文使用 PyTorch 实现。这不是偶然的。
历史原因：TensorFlow 1.x 的"作死"
TensorFlow 1.x 用的是静态图（后面会细讲）。你需要先定义整个计算图，然后在一个 Session 里运行。调试时你甚至不能直接 print 中间变量的值——你必须用 tf.Print。
这是 TensorFlow 1.x 写一个简单的全连接网络：
# TensorFlow 1.x 风格
import tensorflow as tf
x = tf.placeholder(tf.float32, [None, 784])
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
y = tf.nn.softmax(tf.matmul(x, W) + b)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
result = sess.run(y, feed_dict={x: data})
而 PyTorch 的做法：
# PyTorch 风格 —— 和普通 Python 一模一样
x = torch.randn(100, 784)
W = torch.randn(784, 10)
b = torch.randn(10)
y = torch.softmax(x @ W + b, dim=1)
学术界的选择
2017 年 PyTorch 发布后，研究人员发现用它做实验比 TensorFlow 快太多了。调试能直接用 print、能用 pdb 断点、能写任意 Python 控制流——这些在 TensorFlow 1.x 里都是奢望。于是学术界迅速倒向 PyTorch。
而工业界跟着学术界走：招聘时候选人大多熟悉 PyTorch，公司也就顺势切换了。
Google 的反应
Google 在 TensorFlow 2.0 中完全转向了 Eager Execution（动态图），API 设计也大量参考 PyTorch，但为时已晚，心智已经流失。
关键时间线
| 时间 | 事件 |
|---|---|
| 2015.11 | TensorFlow 1.0 发布 |
| 2016.10 | PyTorch 0.1 发布 |
| 2018.12 | PyTorch 1.0 发布（JIT 支持） |
| 2019.10 | TensorFlow 2.0 发布（全面转向动态图） |
| 2020-2023 | PyTorch 市占率超过 80% |
| 2023 | PyTorch 2.0 发布（torch.compile） |
三、Tensor（张量）彻底讲透
小白版：
"张量"这个名字听起来吓人，其实它就是多维数组——Excel 表格的加强版。一个数（标量）是 0 维张量，就是一格；一行数字（向量）是 1 维张量，就是 Excel 的一行；一个表格（矩阵）是 2 维张量；如果再叠多层表格，就是 3 维、4 维的张量。比如一张彩色图片，可以看成 3 张表格叠在一起（红、绿、蓝），那就是 3 维张量。PyTorch 的 Tensor 本质和 NumPy 的多维数组几乎一样，唯一的区别是：Tensor 能自动计算梯度（回想一下高中数学的求导），而且能搬到显卡（GPU）上跑，速度快几十倍。
专业版：
3.1 张量就是多维数组
不要被"张量"这个词吓到。它就是个多维数组，和多维的 NumPy ndarray 本质一样：
标量（0 维）： torch.tensor(3.14) # 形状 []
向量（1 维）： torch.tensor([1, 2, 3]) # 形状 [3]
矩阵（2 维）： torch.randn(3, 4) # 形状 [3, 4]
三维张量： torch.randn(2, 3, 4) # 形状 [2, 3, 4] （比如 2 张 3×4 的灰度图）
四维张量： torch.randn(16, 3, 224, 224) # [batch, channel, height, width]
| 维度 | 实际含义 | 常见场景 |
|---|---|---|
| 0 | 标量 | 单个数值，如 loss 值 |
| 1 | 向量 | 一个样本的特征向量 |
| 2 | 矩阵 | 一批特征、一张灰度图 |
| 3 | 3D 张量 | 一张 RGB 图像（C×H×W）、NLP 中一个 batch 的序列 |
| 4 | 4D 张量 | 一个 batch 的 RGB 图像（N×C×H×W） |
| 5 | 5D 张量 | 视频数据（N×C×T×H×W） |
3.2 PyTorch Tensor vs NumPy ndarray
相同点：
- 切片、索引、广播机制几乎一样
- 大量同名函数：
reshape、sum、mean、argmax等
关键区别：
| PyTorch Tensor | NumPy ndarray | |
|---|---|---|
| GPU 支持 | 原生支持 .cuda() |
不支持 |
| 自动求导 | requires_grad=True 即可 |
不支持 |
| 内存布局 | 行优先（C-contiguous） | 行优先 |
| 底层存储 | 共享存储视图 | 共享存储视图 |
从 NumPy 迁移的最快方式：
import numpy as np
import torch
# NumPy → Tensor
np_arr = np.array([1, 2, 3])
tensor = torch.from_numpy(np_arr)
# Tensor → NumPy（注意：必须在 CPU 上）
tensor_cpu = tensor.cpu()
np_arr_back = tensor_cpu.numpy()
# 直接创建 Tensor（推荐写法）
x = torch.tensor([1, 2, 3]) # 从数据创建
x = torch.zeros(3, 4) # 全零
x = torch.ones(3, 4) # 全一
x = torch.randn(3, 4) # 正态分布随机
x = torch.arange(0, 10, 2) # 类似 range
3.3 三个最容易被忽视的关键属性
x = torch.randn(2, 3, 4)
# 1. shape / size —— 形状
print(x.shape) # torch.Size([2, 3, 4])
print(x.size(0)) # 2
# 2. dtype —— 数据类型
print(x.dtype) # torch.float32
x_int = x.long() # 转为 int64
x_half = x.half() # 转为 float16（省显存但精度低）
# 3. device —— 在哪个设备上
print(x.device) # cpu
x_gpu = x.cuda() # 移到 GPU
x_cpu = x_gpu.cpu() # 移回 CPU
3.4 常见误区：view vs reshape vs permute
误区：以为它们可以互换。
x = torch.randn(2, 3, 4) # shape: [2, 3, 4]
# view：要求内存连续，否则报错；返回视图（共享数据）
x_view = x.view(2, 12) # [2, 12]，OK
# x.transpose(0,1).view(2,12) # RuntimeError! 不连续内存
# reshape：不要求连续，会自动拷贝；更安全但可能额外开销
x_reshaped = x.reshape(2, 12) # 总是 OK
# permute：交换维度顺序，不改变数据
x_permuted = x.permute(2, 0, 1) # [4, 2, 3]
建议： 开发阶段用 reshape，性能优化阶段再考虑 view 和 contiguous。
3.5 广播机制（Broadcasting）
和 NumPy 规则一致：从最后一维开始比对，维度为 1 或缺失的维度会自动扩展。
a = torch.randn(3, 1) # [3, 1]
b = torch.randn(1, 4) # [1, 4]
c = a + b # [3, 4] 自动广播
实际工作中的典型场景：
# 给一个 batch 的图片做归一化
images = torch.randn(16, 3, 224, 224) # [B, C, H, W]
mean = torch.tensor([0.485, 0.456, 0.406]) # [3]
std = torch.tensor([0.229, 0.224, 0.225]) # [3]
# mean 和 std 会自动广播： [3] → [1, 3, 1, 1] → 匹配 [16, 3, 224, 224]
images_normalized = (images - mean[:, None, None]) / std[:, None, None]
四、自动求导（Autograd）真正原理
小白版：
训练神经网络就是不断微调参数让模型更准，而"往哪个方向调、调多少"需要靠求导（梯度）来指导。就像你在山上蒙着眼找最低点，每一步都用脚试探地面是往上还是往下倾斜，然后往下的方向走。PyTorch 的 Autograd 就是自动帮你做这件事：你只写"向前算"的代码（比如输入×参数=输出），它会在后台悄悄记录每一步运算，等你算完说一句 backward()，它就自动倒着走一遍，算出每个参数该怎么调。不用你手写任何求导公式。
专业版：
4.1 为什么需要自动求导
深度学习训练的本质是：不断调整模型参数，使得损失函数值变小。
调整参数 = 损失函数对每个参数求偏导数，然后用梯度下降更新。
一个简单的全连接网络可能有几百万个参数，手算偏导数完全不可能。Autograd 就是解决这个问题的。
4.2 计算图（Computation Graph）
PyTorch 在你执行运算时，会在后台悄悄构建一张有向无环图（DAG），每个节点记录了三样东西：
- 执行了什么操作（加法、乘法、ReLU 等）
- 输入是什么
- 输出是什么
当你调用 .backward() 时，PyTorch 从最终输出反向遍历这张图，用链式法则自动计算每个参数的梯度。
import torch
# 开启梯度追踪
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1
# 此时 y = 4 + 6 + 1 = 11
# dy/dx = 2x + 3 = 7
y.backward() # 反向传播
print(x.grad) # tensor(7.0) ← 自动算出来了
这张图长什么样？
x ──→ pow(2) ──→ add ──→ add ──→ y
↑
x ──→ mul(3) ────────┘
↑
1 ─────────────┘
4.3 requires_grad 的三个真相
# 1. 叶子节点默认不追踪
x = torch.tensor([1.0, 2.0])
print(x.requires_grad) # False
# 需要显式开启
x = torch.tensor([1.0, 2.0], requires_grad=True)
# 2. 权重的 requires_grad 是自动开的
linear = nn.Linear(10, 5)
print(linear.weight.requires_grad) # True
print(linear.bias.requires_grad) # True
# 3. 推理时关掉，省显存、加速
with torch.no_grad():
output = model(input) # 不构建计算图，不追踪梯度
4.4 backward() 干了什么
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum() # y = 1 + 4 + 9 = 14
y.backward()
print(x.grad) # tensor([2., 4., 6.])
# dy/dx_i = 2 * x_i
关键点：
.backward()只能对标量调用（或者传一个与输出同形状的 gradient 参数）- 梯度是累加的——多次调用
.backward()会在.grad上累加 - 这就是为什么每个训练步必须调用
optimizer.zero_grad()
# 如果不 zero_grad 会怎样？
x = torch.tensor([1.0], requires_grad=True)
for _ in range(3):
y = x ** 2
y.backward()
print(x.grad) # 2, 4, 6 —— 梯度累加了！
4.5 detach() 的应用场景
当你需要拿一个张量的值做其他计算，但不想让梯度回传时：
# GAN 训练中：更新 Generator 时冻结 Discriminator
fake_output = generator(noise)
score = discriminator(fake_output.detach()) # detach: 梯度在此截断
detach() 返回一个共享数据但不共享计算图的新张量。
五、动态图 vs 静态图（为什么 PyTorch 赢了）
小白版：
“动态图"和"静态图"的区别，很像"即兴演讲"和"提前写稿”——或者更接地气地说，像"边打游戏边看血条"和"打完一局才能看回放"。
静态图（TensorFlow 1.x） 要求你先把整个计算流程画成一张大图纸，交给"执行引擎"去跑。跑的过程中你不能中途打断、不能看中间结果、不能用普通 if/for。就像你给包工头一份完整的施工蓝图，他在封闭的工地里施工，你想看看厕所的水管走位对不对，对不起，进不去。
动态图（PyTorch） 是你写一行，它就执行一行，你想 print 就 print，想 if 就 if，想 break 就 break。就像你一边写代码一边运行 Jupyter Notebook，写错了马上看到报错，立刻改。这在做 NLP（文本长度不固定）时优势巨大，因为长句子走长路径、短句子走短路径，不用为了对齐浪费算力。
这就是为什么研究人员爱死 PyTorch——做实验、调模型时，你需要随时看到中间结果，而不是等图跑完了才发现写错了。
专业版：
这是理解 PyTorch 最核心优势的一章。
5.1 静态图（Define-and-Run）
代表： TensorFlow 1.x
先定义整个计算图，然后编译，最后运行。就像：
- 你画好整栋楼的图纸
- 交给包工头
- 包工头一次性盖完
# TensorFlow 1.x
graph = tf.Graph()
with graph.as_default():
x = tf.placeholder(tf.float32, shape=[None, 784])
W = tf.Variable(tf.zeros([784, 10]))
y = tf.matmul(x, W)
with tf.Session(graph=graph) as sess:
sess.run(tf.global_variables_initializer())
result = sess.run(y, feed_dict={x: my_data})
静态图的痛苦：
- 不能在定义时打印中间值
- 不能用 Python 原生 if/for（要用
tf.while_loop、tf.cond） - 调试极其困难——出错信息指向的是图里面的某个 op，不是你写的代码行
5.2 动态图（Define-by-Run）
代表： PyTorch
代码执行到哪，图就构建到哪。就像：
- 你一边写代码
- 图一边生成
- 随时可以
print、if、for
# PyTorch：动态图
x = torch.randn(32, 784)
# 可以直接 print
print(f"输入形状: {x.shape}")
# 可以用任意 Python 控制流
if x.shape[0] > 16:
x = x[:16]
# 可以用 for 循环
for i, sample in enumerate(x):
if i == 5:
break
5.3 动态图在 NLP 中的碾压优势
NLP 任务中序列长度是不固定的。静态图里你需要 pad 到最大长度，浪费大量计算。动态图里：
# PyTorch：每句话按实际长度走计算图
sentences = ["Hello", "This is a longer sentence", "Hi"]
for sent in sentences:
tokens = tokenizer(sent)
output = model(torch.tensor(tokens)) # 每条独立走图
# 短句走的图短，长句走的图长，不浪费计算
在静态图中，一个 batch 里所有句子必须 pad 成同样长度，然后在计算时用 mask 把 padding 位置 mask 掉——凭空多了大量无效计算。
5.4 静态图的唯一优势（现在也被追上了）
静态图唯一的优势是：图在运行前已知全部结构，可以做图优化（算子融合、内存复用）。
但 PyTorch 2.0 引入的 torch.compile() 已经用 JIT 技术在运行时做了这些优化，动态图的性能差距几乎消失。
# PyTorch 2.0+
model = MyModel()
model = torch.compile(model) # 一行代码获得静态图级别的优化
六、一个完整模型训练流程（重点）
小白版：
训练一个 AI 模型的流程，和教你家的狗学"坐下"几乎一模一样：
- 准备数据：拿出狗粮（训练图片+标签）
- 给指令：说"坐下"（输入图片给模型）
- 看结果：狗做了某个动作（模型输出预测）
- 给反馈：它做对了奖励、做错了纠正（计算损失函数——“差了多少”）
- 反向调整：狗下次会更倾向做对的动作（反向传播 + 优化器更新参数）
- 重复 1-5：一遍又一遍，直到狗（模型）学会
下面这套代码就是你"训狗"的完整自动化流程。
专业版：
这是你工作中 90% 时间在写的东西。我们用 MNIST 分类来演示。
6.1 完整代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# ========== 1. 数据准备 ==========
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) # MNIST 的经验值
])
train_dataset = datasets.MNIST(
root='./data', train=True, transform=transform, download=True
)
test_dataset = datasets.MNIST(
root='./data', train=False, transform=transform, download=True
)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
# ========== 2. 定义模型 ==========
class MNISTModel(nn.Module):
def __init__(self):
super().__init__()
self.flatten = nn.Flatten()
self.fc1 = nn.Linear(28 * 28, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 10)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.2)
def forward(self, x):
x = self.flatten(x)
x = self.relu(self.fc1(x))
x = self.dropout(x)
x = self.relu(self.fc2(x))
x = self.dropout(x)
x = self.fc3(x)
return x
model = MNISTModel()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
# ========== 3. 损失函数和优化器 ==========
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# ========== 4. 训练循环 ==========
epochs = 10
for epoch in range(epochs):
# ---- 训练阶段 ----
model.train()
train_loss = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
# 前向
output = model(data)
loss = criterion(output, target)
# 反向
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item()
avg_train_loss = train_loss / len(train_loader)
# ---- 验证阶段 ----
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
_, predicted = torch.max(output, 1)
total += target.size(0)
correct += (predicted == target).sum().item()
accuracy = 100.0 * correct / total
print(f"Epoch {epoch+1:2d} | "
f"Train Loss: {avg_train_loss:.4f} | "
f"Test Accuracy: {accuracy:.2f}%")
print("训练完成！")
6.2 每个步骤在做什么
| 步骤 | 代码关键行 | 目的 |
|---|---|---|
| 数据准备 | DataLoader(train_dataset, batch_size=64) |
批量加载数据 |
| 定义模型 | class MNISTModel(nn.Module) |
定义网络结构 |
| 损失函数 | nn.CrossEntropyLoss() |
衡量预测和真实的差距 |
| 优化器 | optim.Adam(model.parameters()) |
决定如何更新参数 |
| 前向传播 | output = model(data) |
计算预测值 |
| 计算损失 | loss = criterion(output, target) |
算出差多少 |
| 清空梯度 | optimizer.zero_grad() |
防止梯度累加 |
| 反向传播 | loss.backward() |
计算每个参数的梯度 |
| 更新参数 | optimizer.step() |
沿梯度方向微调参数 |
| 验证 | model.eval() + torch.no_grad() |
评估模型，不更新参数 |
6.3 model.train() vs model.eval()
这可能是最容易被忽视但影响非常大的一行代码。
model.train() # 训练模式
model.eval() # 评估模式
它们影响哪些层：
| 层 | train() 行为 |
eval() 行为 |
|---|---|---|
| Dropout | 随机丢弃神经元 | 关闭丢弃，所有神经元都工作 |
| BatchNorm | 用当前 batch 的均值和方差 | 用训练期间积累的全局均值和方差 |
如果忘了切换模式：在测试时 Dropout 还在随机丢弃，导致每次推理结果不一样，准确率偏低；BatchNorm 用测试 batch 的统计量而非全局统计量，对结果稳定性也有影响。
6.4 常见误区：Loss 不下降
看到 Loss 不下降时，按这个顺序排查：
# 1. 确认数据进去了
print(data.shape, target.shape)
# 2. 确认模型输出了
output = model(data)
print(output.shape)
# 3. 检查 loss 是否在合理范围
# CrossEntropyLoss 初始 loss 应该在 -ln(1/num_classes) 附近
# 对 10 分类，约 2.3；如果看到天文数字或 NaN，说明有问题
import math
expected_init_loss = -math.log(1 / 10)
print(f"期望初始 loss: {expected_init_loss:.2f}")
# 4. 梯度是否在流动
for name, param in model.named_parameters():
if param.grad is not None:
print(f"{name}: grad_mean={param.grad.mean():.6f}")
七、Dataset 和 DataLoader（工作中必会）
小白版：
训练 AI 需要的数据量往往是十万甚至百万级别，你绝不可能一次性把所有数据塞进内存（电脑会直接卡死）。Dataset 和 DataLoader 解决的就是"怎么高效喂数据"的问题。
打个比方：你在厨房做菜，冰箱（硬盘）里有 1000 个鸡蛋。你不可能一次全拿出来打在碗里（内存爆炸）。你需要的是：每次拿 8 个鸡蛋（batch），打散，下锅，然后再拿下一批。Dataset 就是你的冰箱管理员——它知道每个鸡蛋在哪、什么品质（标签）。DataLoader 就是你的帮手——它帮你按批次取蛋、打蛋（预处理）、送到灶台（GPU）。
专业版：
7.1 为什么需要 Dataset / DataLoader
你不可能一次性把所有数据加载到内存（ImageNet 有 140 万张图）。需要：
- 按需加载：只加载当前 batch 的数据
- 并行加载：用多进程加速
- 随机打乱 + 批处理 + 预处理
这就是 Dataset 和 DataLoader 的职责。
7.2 自定义 Dataset
工作中你面对的不是 MNIST，而是一堆散落文件。自定义 Dataset 是你最常用的技能：
from torch.utils.data import Dataset
from PIL import Image
import os
class MyImageDataset(Dataset):
def __init__(self, image_dir, label_file, transform=None):
self.image_dir = image_dir
self.transform = transform
# 读取标签文件：每行 "image_name.jpg,label"
self.samples = []
with open(label_file, 'r') as f:
for line in f:
img_name, label = line.strip().split(',')
self.samples.append((img_name, int(label)))
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
img_name, label = self.samples[idx]
img_path = os.path.join(self.image_dir, img_name)
# 用 PIL 读取图片
image = Image.open(img_path).convert('RGB')
if self.transform:
image = self.transform(image)
return image, label
三个必须实现的方法：
| 方法 | 作用 |
|---|---|
__init__ |
初始化文件列表、路径、transform |
__len__ |
返回数据集大小 |
__getitem__ |
根据索引返回一个 (data, label) 元组 |
7.3 DataLoader 的关键参数
loader = DataLoader(
dataset,
batch_size=32, # 每个 batch 多大
shuffle=True, # 是否打乱
num_workers=4, # 用几个子进程加载数据 ★ 重要
pin_memory=True, # 锁页内存，加速 CPU→GPU 传输
drop_last=False, # 最后一个不完整 batch 是否丢弃
prefetch_factor=2, # 每个 worker 提前加载的 batch 数
persistent_workers=True # 不 kill worker，避免反复创建进程开销
)
num_workers 怎么调：
# 经验公式（不是绝对的）
# num_workers = min(CPU 核心数, batch_size, 8)
# 太小：数据加载成为瓶颈，GPU 空转
# 太大：CPU 吃不消，多进程创建开销反而拖慢
# 建议：从 4 开始试，逐步增大，找到 GPU 利用率最高且 CPU 不过载的值
pin_memory=True 为什么能加速：
正常流程：数据在可分页内存 → GPU 显存（需要先拷到一个锁页内存缓冲区）
开启后：数据直接加载到锁页内存 → 一步拷到 GPU 显存
速度提升通常在 1.5x ~ 2x。
7.4 collate_fn 的真正用法
当你每个样本不是固定形状（NLP 变长序列、目标检测不同数量框）时，需要自定义 collate_fn：
def my_collate_fn(batch):
"""batch 是一个 list，每个元素是 Dataset.__getitem__ 返回的元组"""
images = [item[0] for item in batch]
labels = [item[1] for item in batch]
# pad 到相同大小
max_h = max(img.shape[1] for img in images)
max_w = max(img.shape[2] for img in images)
padded = []
for img in images:
pad_h = max_h - img.shape[1]
pad_w = max_w - img.shape[2]
padded.append(torch.nn.functional.pad(img, (0, pad_w, 0, pad_h)))
return torch.stack(padded), torch.tensor(labels)
loader = DataLoader(dataset, batch_size=16, collate_fn=my_collate_fn)
八、GPU 与 CUDA 加速
小白版：
一句话解释 CPU 和 GPU 的区别：CPU 是一个博士后，GPU 是一千个小学生。
博士后（CPU）智商极高，能处理复杂的逻辑判断、分支跳转，但人数少（8-16个核心）。一千个小学生（GPU）单个人不聪明，只能做简单的加减乘除，但胜在人多、手快。
训练神经网络的本质就是做海量的矩阵乘法——每个元素都是先乘再加，计算方式一模一样。这正好是"一千个小学生"最擅长的事：分一人一列，同时算。而 PyTorch 做的事情就是帮你把矩阵乘法自动分配给 GPU 上的上千个核心并行计算，所以你写 .cuda() 一行代码就能获得几十倍加速。
专业版：
8.1 GPU 为什么比 CPU 快
CPU 是"全能运动员"：核心少（8-16 个），每个核心能力强，擅长串行、复杂逻辑。
GPU 是"万名小学生"：上千个核心，每个核心简单，擅长做同一件事做很多遍。
矩阵乘法就是"同一件事做很多遍"的典型——A @ B 中每个元素的计算方式完全一样。这正是 GPU 最擅长的。
CUDA 是 NVIDIA 提供的编程接口，让你能直接控制 GPU 做并行计算。PyTorch 在底层调用 CUDA，你只需要 .cuda()。
8.2 多 GPU 的基础认知
误区：“我有两块 GPU，训练速度就翻倍。”
不是。多 GPU 有通信开销，实际加速比：
| GPU 数量 | 理想加速比 | 实际加速比（典型） |
|---|---|---|
| 1 | 1x | 1x |
| 2 | 2x | 1.7x ~ 1.9x |
| 4 | 4x | 2.8x ~ 3.5x |
| 8 | 8x | 4.5x ~ 6.5x |
GPU 越多，通信占比越大，效率越低。
8.3 单卡训练模板
# 最佳实践：自动检测 GPU 可用性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")
# 模型和数据都要移到同一设备
model = MyModel().to(device)
for data, target in loader:
data, target = data.to(device), target.to(device)
# 训练...
常见坑：
# 错误示范 1：模型在 GPU，数据在 CPU
model = MyModel().cuda()
for data, _ in loader:
output = model(data) # RuntimeError! data 还在 CPU 上
# 错误示范 2：张量在不同 GPU 上
x = torch.randn(10).cuda(0)
y = torch.randn(10).cuda(1)
z = x + y # RuntimeError! 不在同一设备
# 正确做法
x = torch.randn(10).cuda(0)
y = torch.randn(10).cuda(0) # 同一设备
z = x + y
九、多 GPU 分布式训练（DataParallel / DDP）
小白版：
如果你有一块 GPU，相当于一个人搬砖。有两块 GPU，你想让两个人一起搬。怎么分活？
DataParallel 方案（不推荐）： 指定一个人当工头（GPU 0），他把一堆砖分成两份，自己拿一份，另一份丢给工友。工友搬完把砖交给工头，工头再统一码好。问题是：工头既要搬砖又要管分配，累死了，而且工友搬完得等工头。
DDP 方案（工业标准）： 两个人各搬各的，搬完之后用对讲机同步一下"我已经搬了 50 块，你呢？""我也 50 块，那我俩平均一下，按 50 块来汇报。"没有工头瓶颈，效率高得多。
这就是为什么生产环境都用 DDP——每个 GPU 平等干活，梯度同步高效，还能跨多台机器。
专业版：
9.1 DataParallel（DP）—— 不推荐但要知道
# 最简单的方式：一行代码
model = nn.DataParallel(model)
# 完整示例
model = MyModel()
if torch.cuda.device_count() > 1:
print(f"使用 {torch.cuda.device_count()} 张 GPU")
model = nn.DataParallel(model)
model = model.cuda()
DP 的工作原理：
- 主卡（GPU 0）拿到一个 batch
- 把 batch 分成 N 份，发给 N 张卡
- 每张卡各自前向 + 反向
- 梯度汇总到主卡
- 主卡更新参数，然后广播给其他卡
DP 的问题：
| 问题 | 影响 |
|---|---|
| 主卡负载不均衡 | GPU 0 显存占用明显高于其他卡 |
| Python GIL 限制 | 多线程受 GIL 影响，效率低 |
| 只支持单机 | 不能跨机器 |
| 无法用混合精度 | 不支持 fp16 |
结论：DP 只适合快速验证，生产环境不要用。
9.2 DistributedDataParallel（DDP）—— 工业标配
DDP 解决了 DP 的所有问题。核心区别：
| DataParallel | DistributedDataParallel | |
|---|---|---|
| 通信方式 | 多线程 + 主卡汇总 | 多进程 + All-Reduce |
| 负载均衡 | 主卡过载 | 每张卡均等 |
| GIL 限制 | 受限制 | 多进程，不受限 |
| 跨机器 | 不支持 | 支持 |
| 混合精度 | 不支持 | 支持 |
| 速度 | 基准 | 比 DP 快 10%-30% |
DDP 的标准写法：
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def setup(rank, world_size):
"""初始化进程组"""
dist.init_process_group(
backend='nccl', # GPU 通信用 NCCL
init_method='env://', # 从环境变量读地址
rank=rank,
world_size=world_size
)
torch.cuda.set_device(rank)
def cleanup():
dist.destroy_process_group()
def train(rank, world_size):
setup(rank, world_size)
# 模型
model = MyModel().to(rank)
model = DDP(model, device_ids=[rank])
# 数据集：每个进程只拿自己的那部分
dataset = MyDataset()
sampler = DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=True
)
loader = DataLoader(
dataset,
batch_size=32,
sampler=sampler, # ★ 关键：用 sampler 而不是 shuffle
num_workers=4,
pin_memory=True
)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(10):
sampler.set_epoch(epoch) # ★ 每个 epoch 重新打乱
for data, target in loader:
data, target = data.to(rank), target.to(rank)
output = model(data)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
cleanup()
# 启动训练
if __name__ == "__main__":
world_size = torch.cuda.device_count()
mp.spawn(train, args=(world_size,), nprocs=world_size)
DDP 的三个关键点：
- 每个 GPU 一个独立进程：
mp.spawn启动多个进程，每个进程管一张卡 - 数据分片：
DistributedSampler确保每张卡拿不同的数据 - 梯度同步：
loss.backward()时自动做 All-Reduce，梯度在所有卡之间平均
9.3 DDP 训练的启动方式
# 单机多卡
torchrun --nproc_per_node=4 train_ddp.py
# 或旧写法（也可以用，但 torchrun 更推荐）
python -m torch.distributed.launch --nproc_per_node=4 train_ddp.py
9.4 单卡训练 vs DP vs DDP 速度对比（典型 ResNet-50）
| 配置 | 训练时间（ImageNet 1 epoch） | 加速比 |
|---|---|---|
| 1×V100 | 45 分钟 | 1.0x |
| 4×V100 DP | 18 分钟 | 2.5x |
| 4×V100 DDP | 14 分钟 | 3.2x |
十、PyTorch Lightning 工程化训练
小白版：
原生 PyTorch 写一个完整的训练流程，100 行代码里 80 行是重复的模板——每次都要手写 zero_grad()、backward()、step()、model.train() 这一套。PyTorch Lightning 就是把这些"每次都要写一遍的破事"全帮你做了，你只需要告诉它三件事：① 数据怎么来 ② 模型长啥样 ③ 训练一步该干嘛。然后它自动处理 GPU 分配、进度条、日志、断点续训、混合精度这些工程细节。
可以理解为：原生 PyTorch 是手动挡汽车，你享受驾驶乐趣但操作繁琐。Lightning 是自动挡，你踩油门（写核心逻辑）就行，不用操心离合和换挡（工程细节）。
专业版：
10.1 原生 PyTorch 的痛点
上面的训练代码写了 100 行，但其中 80 行是样板代码：
optimizer.zero_grad()loss.backward()optimizer.step()model.train()/model.eval()to(device)- 学习率调度
- checkpoint 保存/加载
- 日志记录
- 16 位精度训练
PyTorch Lightning 把这些全部抽象掉，让你只写模型逻辑。
10.2 用 Lightning 重写 MNIST
import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
class LitMNIST(pl.LightningModule):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28 * 28, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 10)
self.dropout = nn.Dropout(0.2)
def forward(self, x):
x = x.view(x.size(0), -1)
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = F.relu(self.fc2(x))
x = self.dropout(x)
x = self.fc3(x)
return x
# ↓ 你只需要写这一个方法
def training_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.cross_entropy(logits, y)
self.log('train_loss', loss)
return loss
def validation_step(self, batch, batch_idx):
x, y = batch
logits = self(x)
loss = F.cross_entropy(logits, y)
preds = torch.argmax(logits, dim=1)
acc = (preds == y).float().mean()
self.log('val_loss', loss, prog_bar=True)
self.log('val_acc', acc, prog_bar=True)
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=0.001)
# 数据模块
class MNISTDataModule(pl.LightningDataModule):
def __init__(self, batch_size=64):
super().__init__()
self.batch_size = batch_size
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
def setup(self, stage=None):
self.mnist_train = datasets.MNIST(
'./data', train=True, download=True, transform=self.transform
)
self.mnist_val = datasets.MNIST(
'./data', train=False, download=True, transform=self.transform
)
def train_dataloader(self):
return DataLoader(self.mnist_train, batch_size=self.batch_size, shuffle=True)
def val_dataloader(self):
return DataLoader(self.mnist_val, batch_size=self.batch_size)
# 训练，就 3 行
model = LitMNIST()
dm = MNISTDataModule()
trainer = pl.Trainer(max_epochs=10, accelerator='auto', devices='auto')
trainer.fit(model, dm)
对比： 原生 PyTorch 的 100 多行变成了约 60 行，而且：
- 自动处理
model.train()/model.eval() - 自动处理 GPU
- 自动打印进度条和日志
- 自动保存最优 checkpoint
- 支持多 GPU / TPU / 混合精度——改一行配置即可
10.3 Lightning 的进阶配置
trainer = pl.Trainer(
max_epochs=100,
# 硬件
accelerator='gpu', # 'cpu', 'gpu', 'tpu', 'auto'
devices=4, # GPU 数量
strategy='ddp', # 'dp', 'ddp', 'fsdp', 'deepspeed'
# 精度
precision='16-mixed', # 混合精度训练，省显存、加速
# 日志
logger=pl.loggers.TensorBoardLogger('logs/'),
log_every_n_steps=10,
# checkpoint
callbacks=[
pl.callbacks.ModelCheckpoint(
monitor='val_acc',
mode='max',
save_top_k=3,
filename='mnist-{epoch:02d}-{val_acc:.2f}'
),
pl.callbacks.EarlyStopping(
monitor='val_loss',
patience=5,
mode='min'
),
pl.callbacks.LearningRateMonitor(logging_interval='epoch')
]
)
十一、真实 AI 项目训练流程
小白版：
拿到一个真实 AI 任务，不要上来就写代码。遵循四步走：
- 先跑通：用最简单的模型、最小量的数据，先把训练→验证→保存整个流程跑起来，确认没 bug。
- 再调对：把小数据上的 loss 压到接近 0——证明模型有能力学会。
- 再调好：在全量数据上，逐个调整学习率、batch size、数据增强等参数，让验证指标达标。
- 最后工程化：加日志、加断点续训、加模型导出、加推理接口。
这和做菜一样：先确认锅和灶都能用（跑通），再确认调料比例没错（调对），再反复微调火候（调好），最后装盘上菜（工程化）。
专业版：
阶段一：项目启动（1-2 天）
1. 明确任务：分类 / 检测 / 分割 / 生成？
2. 确定评价指标：准确率 / mAP / IoU / BLEU？
3. 找基线模型（Baseline）：github / Papers with Code
4. 搭建最小可运行环境
# 标准项目结构
project/
├── configs/ # 配置文件（推荐用 yaml）
│ └── baseline.yaml
├── data/ # 原始数据
│ ├── train/
│ └── val/
├── dataset/ # Dataset 类
│ └── my_dataset.py
├── models/ # 模型定义
│ └── my_model.py
├── utils/ # 工具函数
│ ├── metrics.py
│ └── visualize.py
├── train.py # 训练入口
├── inference.py # 推理入口
└── requirements.txt
阶段二：快速验证（1 天）
目标：用最简单的方法跑通全流程，确认代码能收敛。
# 先在小数据上跑通
tiny_dataset = torch.utils.data.Subset(full_dataset, range(1000))
# 过拟合小数据——确认模型能学习
model = MyModel()
optimizer = optim.Adam(model.parameters(), lr=0.01)
for epoch in range(100):
# ... 训练，不调参，纯粹看 loss 能不能降到接近 0
pass
# loss 降到接近 0 → 模型没问题
# loss 降不下去 → 代码有 bug 或模型容量不够
阶段三：调参优化（3-7 天）
关键超参数和调参顺序（按重要性排序）：
| 优先级 | 超参数 | 典型范围 | 对效果影响 |
|---|---|---|---|
| 1 | 学习率 | 1e-5 ~ 1e-2 | 最大 |
| 2 | batch size | 16 / 32 / 64 / 128 | 大 |
| 3 | 优化器 | Adam / AdamW / SGD + Momentum | 中 |
| 4 | 正则化 | Dropout 0.1~0.5, Weight Decay 0.0001~0.01 | 中 |
| 5 | 学习率调度 | Cosine / Step / ReduceLROnPlateau | 中 |
| 6 | 数据增强 | Rotate / Flip / Cutout / MixUp | 大 |
# 学习率调度示例
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
# 余弦退火：从初始 lr 平滑降到 0
scheduler = CosineAnnealingLR(optimizer, T_max=epochs)
# 或：监控 val_loss，不降了就减半
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3)
for epoch in range(epochs):
train()
val_loss = validate()
scheduler.step(val_loss) # ReduceLROnPlateau 需要传 val_loss
# scheduler.step() # CosineAnnealing 不需要参数
阶段四：完整训练与产出
# 保存完整训练状态（不只是权重）
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
'config': config_dict,
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pt')
# 推理时只加载权重
model = MyModel()
model.load_state_dict(torch.load('best_model.pt')['model_state_dict'])
model.eval()
十二、常见坑与避坑指南
小白版：
PyTorch 初学者踩坑最多的地方可以总结为"三板斧"：
- 忘清零：每次训练前要
optimizer.zero_grad()，否则梯度会像滚雪球一样越滚越大，参数更新越来越离谱。 - 模式忘切换：训练时用
model.train()，测试时用model.eval()。忘了切换的话，测试时 Dropout 还在随机丢神经元、BatchNorm 还在用错误统计量，准确率直接掉坑里。 - 数据不在同一设备：模型搬到了 GPU 上，但输入数据还在 CPU 上——程序直接报错。记住一句话：模型和数据必须待在同一张卡（或 CPU）上。
除此之外还有一些更隐蔽的坑，下面详细拆解。
专业版：
坑 1：忘了 zero_grad()
# ❌ 错误
for data, target in loader:
loss = criterion(model(data), target)
loss.backward()
optimizer.step() # 梯度一直累加，参数更新量越来越大
# ✅ 正确
for data, target in loader:
loss = criterion(model(data), target)
optimizer.zero_grad() # 每次迭代前清空梯度
loss.backward()
optimizer.step()
坑 2：在 with torch.no_grad() 外面做验证
# ❌ 错误：验证时也在构建计算图，浪费显存
model.eval()
for data, target in val_loader:
output = model(data)
loss = criterion(output, target) # 计算图被构建了！
# ✅ 正确
model.eval()
with torch.no_grad():
for data, target in val_loader:
output = model(data)
loss = criterion(output, target) # 不构建计算图
坑 3：在 eval() 模式下忘了关掉，又去训练
# ❌ 灾难性错误
model.eval()
# ... 做了验证
# 忘了 model.train()
for data, target in train_loader:
output = model(data) # Dropout 关闭了，BatchNorm 用全局统计量
# 模型训练效果大打折扣
坑 4：nn.Module 中用了 Python list 存子模块
# ❌ 错误：list 里的模块不会被注册
class BadModel(nn.Module):
def __init__(self):
super().__init__()
self.layers = [nn.Linear(10, 10) for _ in range(5)]
def forward(self, x):
for layer in self.layers:
x = layer(x)
return x
model = BadModel()
print(list(model.parameters())) # 空的！
# ✅ 正确：用 nn.ModuleList
class GoodModel(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(5)])
def forward(self, x):
for layer in self.layers:
x = layer(x)
return x
model = GoodModel()
print(len(list(model.parameters()))) # 10（5个weight + 5个bias）
坑 5：torch.load 的安全风险
# ❌ 危险：从不可信来源加载模型
model = torch.load('downloaded_from_internet.pt')
# 这个 .pt 文件可能包含恶意 pickle 代码，直接执行！
# ✅ 安全做法
model = MyModel()
checkpoint = torch.load('model.pt', weights_only=True)
model.load_state_dict(checkpoint['model_state_dict'])
坑 6：混合 tensor 类型
# ❌ 常见错误
x = torch.randn(10) # float32
y = torch.tensor([1, 2, 3, 4]) # int64
z = x + y # RuntimeError!
# ✅ 统一类型
y = torch.tensor([1, 2, 3, 4], dtype=torch.float32)
z = x + y # OK
坑 7：显存泄漏
# ❌ 每个 batch 都在创建新张量但不释放
losses = []
for data, target in loader:
output = model(data)
loss = criterion(output, target)
losses.append(loss.item()) # .item() 没问题，但如果不加 .item()
# loss 是个标量张量，如果放进 list，引用了计算图
# 计算图越来越大，显存爆炸
# ✅ 正确
losses = []
for data, target in loader:
output = model(data)
loss = criterion(output, target)
losses.append(loss.item()) # 转为 Python 数字，切断引用
loss.backward()
optimizer.step()
optimizer.zero_grad()
十三、一句话总结 + 学习路线
小白版：
PyTorch 就是"深度学习的 Python"——它让你用写普通 Python 的方式写 AI。你不需要理解显卡底层怎么工作、不需要手算梯度、不需要自己实现矩阵乘法。你只需要学会三件事：定义模型（搭积木）、加载数据（喂数据）、写训练循环（调参数），就能做出能用的 AI。
推荐从小练习开始：先用 PyTorch 实现一个线性回归（10 行代码），再做 MNIST 手写数字识别（50 行代码），然后逐步进阶。不要只看不写——每个阶段写一个小项目，代码量是学会的唯一标准。
专业版：
一句话总结
PyTorch 就是把 NumPy 搬上 GPU，加上自动求导，然后用 Python 原生风格写深度学习代码的框架——它的成功来自于"让写深度学习像写普通 Python 一样自然"。
推荐学习路线
第一阶段（1 周）
├── Tensor 基本操作（创建、索引、运算、广播）
├── Autograd（requires_grad, backward, no_grad）
└── 小练习：用纯 Tensor + Autograd 实现线性回归
第二阶段（1 周）
├── nn.Module（定义模型、forward 方法）
├── nn.Linear, nn.Conv2d, nn.ReLU 等常用层
├── Loss + Optimizer + 训练循环
└── 小练习：在 MNIST 上训练一个分类器
第三阶段（2 周）
├── Dataset + DataLoader + transforms
├── GPU 训练
├── 保存 / 加载模型
└── 小练习：在自己的数据集上训练 ResNet
第四阶段（1 周）
├── 学习率调度
├── TensorBoard 可视化
├── 多 GPU（DDP）
└── 小练习：用 DDP 训练，打 Kaggle 一个比赛
第五阶段（按需）
├── PyTorch Lightning（工程化）
├── 混合精度训练
├── ONNX / TorchScript 模型部署
├── torch.compile（PyTorch 2.0）
└── 读论文 + 复现
核心资料
| 资源 | 链接 | 适合阶段 |
|---|---|---|
| PyTorch 官方 60 min 教程 | pytorch.org/tutorials | 第一阶段 |
| d2l.ai（动手学深度学习） | d2l.ai | 第二、三阶段 |
| PyTorch Lightning 文档 | lightning.ai/docs | 第四阶段 |
| Papers with Code | paperswithcode.com | 全阶段 |
这篇文章覆盖了从零基础到工业级多 GPU 训练的全部核心知识。前半段用通俗类比建立直觉，后半段用代码和表格给出可实操的方案。如果你正在入门 AI，建议按学习路线一步步来，每学完一个阶段就动手做一个小项目——代码写出来才能真正理解。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。
更多推荐
5
0- 0
已为社区贡献16条内容
所有评论(0)