PyTorch06-PyTorch进阶训练技巧
六、PyTorch进阶训练技巧
6.1 自定义损失函数
尽管PyTorch官方针对主流模型为用户提供了很多内置损失函数,但随着应用场景的变化与自定义模型的出现,那些内置损失函数可能并不适合,于是产生了一些针对非通用模型的损失函数。有些损失函数是针对某特定场景使用,可以直接拿来使用或修改使用,有些场景或自定义模型甚至找不到现成或逼近的损失函数可用,此时就需要自定义损失函数。
本章节学习的内容包括:如何自定义损失函数。
6.1.1 以函数形式定义损失函数
损失函数其实就是用来计算模型输出数据与真实数据之间的差距的,所以计算损失的过程就是一个函数执行的过程。 所以理论上来说,可以通过定义一个函数的形式来自定义损失函数。
1 | |
6.1.2 以类形式定义损失函数
6.1.2.1 类形式定义Loss介绍
虽然以普通函数方式定义损失函数简单、快速,但实践中仍是常以类形式自定义损失。通过查看PyTorch源码可以发现,内置的损失函数的有的继承自torch.nn.modules.loss._Loss,有的继承自torch.nn.modules.loss._WeightedLoss
,甚至还有很多继承自torch.nn.Module,其实就三者之间也存在继承关系,本质这些内置的损失函数仍是继承自torch.nn.Module。
三者之间的继承关系如下。 1
2
3
4
5
6
7
8import torch.nn as nn
# 导入这两个底层抽象基类
from torch.nn.modules.loss import _Loss, _WeightedLoss
# 验证继承关系
print(issubclass(_Loss, nn.Module)) # 输出: True
print(issubclass(_WeightedLoss, _Loss)) # 输出: True
类名开头有下划线_在 Python 中代表这是 PyTorch
内部使用的私有 / 保护基类,官方通常不建议普通用户在自定义 Loss
时直接继承它们,但了解它们对于理解 PyTorch
源码极其有帮助。就是说,_Loss 和 _WeightedLoss
这两个抽象基类是给PyTorhc内置的损失函数使用的,一般不面向用户直接使用,用户要自定义损失函数还是得直接继承自
torch.nn.Module,这也表示我们将此类自定义损失函数当作神经网络的一层来对待。
以下以DiceLoss讲述以类的尬啊自定义损失函数 及使用。 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77# 10-自定义损失函数DiceLoss2.py
import torch
from torch import nn
class DiceLoss(nn.Module):
"""用于二分类或多标签分割的 Dice 损失。
inputs: 模型直接输出的 logits,形状为 (N, C, ...),不应预先经过 sigmoid。
targets: 与 inputs 形状相同、元素为 0 或 1 的分割标签。
每个样本、每个通道分别计算 Dice,再根据 reduction 返回损失,
不会将整个 batch 的样本混在一起计算一个全局 Dice。
"""
def __init__(self, smooth=1e-6, reduction="mean"):
super().__init__()
if reduction not in {"mean", "sum", "none"}:
raise ValueError("reduction 必须是 'mean'、'sum' 或 'none'。")
self.smooth = smooth
self.reduction = reduction
def forward(self, inputs, targets):
if inputs.shape != targets.shape:
raise ValueError(
"inputs 与 targets 的形状必须相同,"
f"当前分别为 {tuple(inputs.shape)} 和 {tuple(targets.shape)}。"
)
if inputs.ndim < 2:
raise ValueError("inputs 与 targets 至少应为 (N, C) 两维张量。")
# 模型应输出 logits;这里统一进行 sigmoid,得到范围为 [0, 1] 的预测概率。
probabilities = torch.sigmoid(inputs)
# 标签通常是整型 0/1 张量,转换为与概率相同的 dtype 和设备后再参与计算。
targets = targets.to(device=inputs.device, dtype=probabilities.dtype)
# 除 batch 维 N 外,沿通道与所有空间维求和。
# 对 (N, C, H, W) 输入, dice 的形状为 (N, C)。
reduce_dims = tuple(range(2, inputs.ndim))
# 默认sum方法的keepdim参数为False,不会保留对应的维度,而是直接消除对应维度即每次减少一个维度
# 从第2维开始消起,直到最后一维。最后结果的形状是: (targets.size(0), targets.size(1))
intersection = (probabilities * targets).sum(dim=reduce_dims)
cardinality = probabilities.sum(dim=reduce_dims) + targets.sum(dim=reduce_dims)
dice = (2 * intersection + self.smooth) / (cardinality + self.smooth)
loss = 1 - dice
if self.reduction == "mean":
return loss.mean()
if self.reduction == "sum":
return loss.sum()
return loss # 形状为 (N, C),可用于自行设置各样本/类别权重。
# ==================== 使用方法:二分类分割 ====================
torch.manual_seed(42) # 使本示例每次运行的随机输入与输出可复现。
# N=2 表示 batch_size;C=1 表示每个像素只有“前景/背景”两个取值。
# inputs 是模型的原始输出 logits,因此使用 torch.randn(),而非范围 [0, 1] 的随机概率。
shape = (2, 1, 4, 4)
inputs = torch.randn(*shape, requires_grad=True)
# 分割标签必须与 logits 形状一致,并且元素为 0 或 1。
targets = torch.randint(0, 2, size=shape)
criterion = DiceLoss()
loss = criterion(inputs, targets)
# DiceLoss 的输出是标量,可直接反向传播;这里仅验证梯度能够正常计算。
loss.backward()
print(f"模拟的预测数据inputs:{inputs}")
print(f"\n模拟的真实数据targets:{targets}")
print(f"\n经过DiceLoss损失函数计算得到的损失值为:{loss}")
print("每个输入元素是否都获得梯度:", inputs.grad is not None)
# 注意:多类别互斥分割(每个像素仅属于 C 个类别之一)通常使用 softmax + one-hot 标签,
# 或将 DiceLoss 与 CrossEntropyLoss 组合;不能直接将本二分类/多标签版本用于该场景。
6.1.2.2 Dice Loss损失函数
Dice 相似系数(Dice Similarity Coefficient,简称 DSC) 是图像分割(如医学图像分割 UNet)领域中最核心的评价指标之一。 基于该系数构造的损失函数被称为 Dice Loss。
6.1.2.2.1 公式及各项符号的含义
公式: \[DSC = \frac{2\vert{}X \cap Y\vert{}}{\vert{}X\vert{} + \vert{}Y\vert{}}\]
- \(X\):通常代表模型预测的分割区域(Prediction Mask)。
- \(Y\):通常代表真实的标签区域(Ground Truth Mask)。
- \(\vert{}X \cap Y\vert{}\)(分子):\(X\) 与 \(Y\) 的交集大小(即模型预测正确且属于真实目标的像素/前景数量)。乘以 \(2\) 是为了进行归一化平衡。
- \(\vert{}X\vert{} + \vert{}Y\vert{}\)(分母):\(X\) 的像素元素总数与 \(Y\) 的像素元素总数直接相加。
- \(DSC\)
的取值范围:在 \(0\)
到 \(1\) 之间。
- \(DSC =
1\):代表预测区域与真实区域完全重合(完美分割)。
- \(DSC = 0\):代表预测区域与真实区域没有任何重合。
- \(DSC =
1\):代表预测区域与真实区域完全重合(完美分割)。
6.1.2.2.2 从 DSC 指标到 Dice Loss
神经网络训练时需要最小化 Loss(损失越小越好),而 DSC 是越大越好。因此,在代码中定义 Dice Loss 时,通常采用以下两种转化方式之一: \[\text{Dice Loss} = 1 - DSC = 1 - \frac{2\vert{}X \cap Y\vert{}}{\vert{}X\vert{} + \vert{}Y\vert{}}\] 或者使用负数形式: \[\text{Dice Loss} = -DSC\] 方式一是最常用形式。这样,当分割效果越好(\(DSC \to 1\))时,\(\text{Dice Loss}\) 就越接近 \(0\)。
6.1.2.2.3 在 PyTorch 代码中如何表达?
在连续值的神经网络输出中(例如 Sigmoid 输出的概率值 \(p \in [0, 1]\)),集合操作转换如下: -
交集 \(\vert{}X \cap
Y\vert{}\) \(\rightarrow\)
逐元素相乘并求和:torch.sum(input * target) -
各自大小 \(\vert{}X\vert{} +
\vert{}Y\vert{}\) \(\rightarrow\)
逐元素求和:torch.sum(input) + torch.sum(target)
并且在分母加上微小的平滑项 epsilon(如
1e-5),防止分母为 0:
1 | |
6.1.2.2.4 为什么分割任务中非常喜欢用 Dice Loss?
普通的交叉熵损失(Cross Entropy Loss)是对所有像素一视同仁地计算误差。如果在医学图像中,病灶(前景目标)只占全图的 1%,背景占了 99%(严重的前背景极度不平衡),模型即使把所有像素都预测为背景,准确率也能达到 99%,但这种模型是无效的。
Dice Loss 关注的是前景与真实区域的交叉重叠比例,它天然消除了背景像素占绝大多数带来的影响,能够极大地改善类别不平衡(Class Imbalance)问题。
6.1.2.3 DiceBCELoss
1 | |
6.1.2.4 IoULoss
1 | |
6.1.2.5 FocalLoss
1 | |
6.1.3 为什么常用“类”的形式
既然定义一个普通函数就行,那什么情况下需要把 Loss 写成继承
nn.Module 的类呢?主要有以下 3 个场景:
损失函数包含状态或超参数(Hyperparameters) 比如之前写的
DiceLoss(smooth=1.0)或CombinedLoss(alpha=0.5),如果写成类,可以在__init__中把smooth或alpha保存为成员变量,调用时只需传criterion(output, target),接口更加统一规范。损失函数包含“可学习参数”(Learnable Parameters) 有些高级损失函数(如某些自适应损失、感知损失、带有可训练权重系数的多任务 Loss),Loss 内部自己也带有需要随训练更新的
nn.Parameter。此时必须使用nn.Module,这样model和loss的参数才能一起被传给优化器(Optimizer):1
optimizer = torch.optim.Adam(list(model.parameters()) + list(criterion.parameters()))兼容 PyTorch 标准生态管道 许多高级训练框架(如 PyTorch Lightning、HuggingFace
Trainer)内部要求criterion必须是一个nn.Module实例,方便统一进行.to(device)设备迁移或数据并行(DataParallel)。
大致的经验就是: - 临时测试、无状态计算、简单数学组合 \(\rightarrow\) 用函数形式
def my_loss(...)(简洁、高效、写起来最快)。 -
模块化封装、包含可调参数/可学习权重、需要对接通用训练框架 \(\rightarrow\) 用继承 nn.Module
的类。
6.2 动态调整学习率
在使用PyTorch编写神经网络模型并进行训练时,如何确定一个合适的学习率或学习率调整策略是一个比较关键的问题。 使用固定的学习率时的话,如果将学习率设置得太小,学习模型参数更新太慢,可能使网络模型在既定总轮次下无法收敛到最佳状态,如果直接增加训练轮次将当于人为增加了训练时间;如果将学习率设置得过大,则模型参数更新过快,可能导致模型无法精准落入收敛区甚至在最优解两侧来回跳动、无法收敛。
为此,人们想到了根据实际情况不断自动更新学习率的学习率调整策略(学习率scheduler)。最常见最简单的学习率调整策略是:在训练初始时将学习率设置恰当大(这大概率需要使用固定的学习率提前训练一轮或几轮,提前摸查下模型的收敛情况),然后随着训练轮次的增加不断减少学习率(比如每20轮将学习率变更原来的1/2)。 当然,也有更新复杂的随实际情况更新学习率的学习率调整策略。PyTorch中内置了很多不同的学习率调整策略;此外,还可以根据需要自定义学习率调整函数或策略。
本节学习的内容包括:根据需要选择与使用PyTorch内置的学习率调整策略、自定义学习率调整策略。
6.2.1 使用官方scheduler
6.2.1.1 查看内置学习率调整策略
在训练神经网络的过程中,学习率是最重要的超参数之一。PyTorch中已经内置大概15个学习率调整策略类,它们放置在包torch.optim.lr_scheduler内,它们都是torch.optim.lr_scheduler.LRScheduler的子类。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18CLASSES
builtins.object
LRScheduler
ChainedScheduler
ConstantLR
CosineAnnealingLR
CosineAnnealingWarmRestarts
CyclicLR
ExponentialLR
LambdaLR
LinearLR
MultiStepLR
MultiplicativeLR
OneCycleLR
PolynomialLR
ReduceLROnPlateau
SequentialLR
StepLR
也可以直接在PyTorch官方文档中查看它们的详细介绍:PyTorch官网optim文档 。每个学习率调整策略类的详细文档如下:
lr_scheduler.LambdaLRlr_scheduler.MultiplicativeLRlr_scheduler.StepLRlr_scheduler.MultiStepLRlr_scheduler.ExponentialLRlr_scheduler.CosineAnnealingLRlr_scheduler.ReduceLROnPlateaulr_scheduler.CyclicLRlr_scheduler.OneCycleLRlr_scheduler.CosineAnnealingWarmRestartslr_scheduler.ConstantLRlr_scheduler.LinearLRlr_scheduler.PolynomialLRlr_scheduler.ChainedSchedulerlr_scheduler.SequentialLR
6.2.1.2 使用内置学习率调整策略
使用学习率调整策略方法较简单,就是在优化器对象的基础上再套一层,将优化器对象作为参数传递给学习率调整策略类创建学习率调整策略实例,最后通过调用学习率调整策略类的step方法,达到调用其内部定义好的逻辑与更新优化器的学习率等。 如果使用固定的学习率,就没有在优化器实例基础上创建与使用的学习率调整策略实例。
可完整运行的实例如下: 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175# 20-使用官方内置策略调整学习率.py
"""使用 PyTorch 官方内置策略动态调整学习率。
本例使用一份可控的二维线性回归数据,不需要下载数据集,直接运行即可。
重点是理解官方 scheduler 的基本使用顺序:
先由 optimizer 更新模型参数,再在一个 epoch 结束时调用 scheduler.step()。
"""
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
# 固定随机种子,使每次运行生成的数据和训练结果基本一致,便于学习和比较。
torch.manual_seed(42)
# 选择运行设备:有可用 GPU 时使用 GPU;否则自动退回到 CPU。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def make_dataset():
"""构造 y = 2*x1 - 3*x2 + 0.5 + noise 形式的回归数据集。"""
# features 的形状为 (200, 2):200 个样本,每个样本有 2 个特征。
features = torch.randn(200, 2)
# labels 的形状为 (200, 1)。保留最后这个维度,便于和模型输出对齐。
noise = 0.1 * torch.randn(200, 1)
labels = 2 * features[:, 0:1] - 3 * features[:, 1:2] + 0.5 + noise
# 前 160 个样本用于训练,后 40 个样本用于验证。
train_dataset = TensorDataset(features[:160], labels[:160])
valid_dataset = TensorDataset(features[160:], labels[160:])
return train_dataset, valid_dataset
def validate(model, data_loader, loss_fn):
"""在验证集上计算平均损失;验证阶段不需要计算梯度。"""
model.eval() # 切换到验证模式;本例没有 BN/Dropout,但这是规范写法。
total_loss = 0.0
total_samples = 0
# torch.no_grad() 会关闭梯度记录,节省验证时的内存和计算。
with torch.no_grad():
for features, labels in data_loader:
features = features.to(device)
labels = labels.to(device)
predictions = model(features)
loss = loss_fn(predictions, labels)
# loss 是当前批次的平均损失,因此乘以批量大小后再累加。
batch_size = features.size(0)
total_loss += loss.item() * batch_size
total_samples += batch_size
return total_loss / total_samples
def train_one_epoch(model, data_loader, loss_fn, optimizer):
"""完成一个 epoch 的训练,并返回该 epoch 的平均训练损失。"""
model.train() # 切换到训练模式。
total_loss = 0.0
total_samples = 0
for features, labels in data_loader:
features = features.to(device)
labels = labels.to(device)
# 1. 清空上一批次留下的梯度。PyTorch 默认会累积梯度。
optimizer.zero_grad()
# 2. 前向传播:predictions 的形状为 (batch_size, 1)。
predictions = model(features)
loss = loss_fn(predictions, labels)
# 3. 反向传播:loss.backward() 计算各参数的梯度,存入 parameter.grad。
loss.backward()
# 4. 根据当前批次的梯度更新模型参数。
# 这一步必须发生在 scheduler.step() 之前。
optimizer.step()
# 记录当前批次的总损失,用于计算整个 epoch 的平均损失。
batch_size = features.size(0)
total_loss += loss.item() * batch_size
total_samples += batch_size
return total_loss / total_samples
def main():
train_dataset, valid_dataset = make_dataset()
# shuffle=True 让每个 epoch 的训练样本顺序不同;验证集不需要打乱。
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
# nn.Linear(2, 1) 接收形状为 (batch_size, 2) 的输入,输出 (batch_size, 1)。
model = nn.Linear(in_features=2, out_features=1).to(device)
loss_fn = nn.MSELoss()
# optimizer 负责根据反向传播得到的梯度更新模型参数。
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
# StepLR 是官方内置的学习率调度器:每经过 step_size 个 epoch,
# 将当前学习率乘以 gamma。这里学习率会是:
# epoch 1~20: 0.1;epoch 21~40: 0.05;epoch 41~60: 0.025,以此类推。
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer, step_size=20, gamma=0.5
)
# 同一段训练通常只选择一种学习率策略,而不应同时创建多个 scheduler
# 一起修改同一个 optimizer 的学习率。下面是几种可替换 StepLR 的官方策略:
#
# scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
# scheduler = torch.optim.lr_scheduler.MultiStepLR(
# optimizer, milestones=[30, 60, 80], gamma=0.1
# )
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
# optimizer, T_max=100, eta_min=0.001
# )
#
# 注意:若改为 ReduceLROnPlateau ,需要在每个 epoch 末尾改为
# scheduler.step(valid_loss),让它根据验证损失是否改善来决定是否降低学习率。
epochs = 100
for epoch in range(epochs):
# 进行训练:train_one_epoch 内部会完成多次 optimizer.step()。
train_loss = train_one_epoch(model, train_loader, loss_fn, optimizer)
valid_loss = validate(model, valid_loader, loss_fn)
# 每个 epoch 的训练(即 optimizer 的参数更新)结束后,再更新学习率。
# 对 StepLR、ExponentialLR、MultiStepLR 等策略,直接调用 scheduler.step()。
# 调用后,下一轮 epoch 会使用更新后的学习率。
scheduler.step()
current_lr = optimizer.param_groups[0]["lr"]
# print(f"optimizer.param_groups:{optimizer.param_groups}")
# 每 10 轮打印一次;最后一轮也会打印,方便观察损失和学习率的变化。
if (epoch + 1) % 10 == 0 or epoch == 0:
print(
f"Epoch [{epoch + 1:3d}/{epochs}] "
f"train_loss: {train_loss:.6f} "
f"valid_loss: {valid_loss:.6f} "
f"next_lr: {current_lr:.5f}"
)
# 理论参数是 weight=[2, -3]、bias=0.5;训练结果应接近它们。
print("\n训练完成")
print("学习到的 weight:", model.weight.detach().cpu())
print("学习到的 bias: ", model.bias.detach().cpu())
if __name__ == "__main__":
main()
"""
### 某次执行时的输出:
(ai_py312) jsy@jmac 06第六章 % python 20-使用官方内置策略调整学习率.py
Epoch [ 1/100] train_loss: 6.632476 valid_loss: 4.458079 next_lr: 0.10000
Epoch [ 10/100] train_loss: 0.056313 valid_loss: 0.033541 next_lr: 0.10000
Epoch [ 20/100] train_loss: 0.009557 valid_loss: 0.009020 next_lr: 0.05000
Epoch [ 30/100] train_loss: 0.009557 valid_loss: 0.009164 next_lr: 0.05000
Epoch [ 40/100] train_loss: 0.009519 valid_loss: 0.009063 next_lr: 0.02500
Epoch [ 50/100] train_loss: 0.009512 valid_loss: 0.008966 next_lr: 0.02500
Epoch [ 60/100] train_loss: 0.009532 valid_loss: 0.008914 next_lr: 0.01250
Epoch [ 70/100] train_loss: 0.009511 valid_loss: 0.008954 next_lr: 0.01250
Epoch [ 80/100] train_loss: 0.009513 valid_loss: 0.008958 next_lr: 0.00625
Epoch [ 90/100] train_loss: 0.009512 valid_loss: 0.009013 next_lr: 0.00625
Epoch [100/100] train_loss: 0.009510 valid_loss: 0.008989 next_lr: 0.00313
训练完成
学习到的 weight: tensor([[ 1.9928, -2.9866]])
学习到的 bias: tensor([0.4892])
"""
6.2.2 自定义scheduler
前面提到过,如果PyTorch中内置的学习率调整策略无法满足实际的需求时,用户可自定义学习率调整方法或学习率调整策略类,以在训练过程中对学习率进行自适应更新。
不论自定义学习率调整方法,还是自定义学习率调整策略类,使用上一般都是先是调用优化器的step方法更新模型权重,再调用自定义学习率调整方法更新学习率或自定义学习率调整策略实例的step方法更新学习率。(标准写法是:先
optimizer.step() 更新模型权重,再
scheduler.step() 更新学习率。这个顺序不是 PyTorch
强制要求的,但它是语义上最合理的写法——保证本轮权重更新用的是本轮开始时的学习率,学习率调度是“跨epoch生效”的。如果是自定义学习率调整方法(函数式),调用时机完全由用户控制,但建议同样遵循“先权重后学习率”的顺序。)
6.2.2.1 自定义学习率调整方法
以下完整示例展示了,使用自定义学习率调整方法在模型训练过程中更新学习率,它在更新模型权重后被调用。
1 | |
6.2.2.2 自定义学习率调整策略类
以下完整示例实现的功能与学习率相关调整策略逻辑与上面的示例完全一样,只是以类的形式实现。自定义学习率调整实例调用step方法更新学习率仍在更新模型权重后被执行。
1 | |
6.3 模型微调torchvision
本节学习的内容包括:模型微调的流程、PyTorch中内置的主要开源模型、如何训练模型的指定层。
6.3.1 微调相关背景介绍
开源社区已经为大家提供了很多针对各种场景下使用的神经网络,其中很多神经网络训练时采用了大规模、高质量的数据集,但它们终归只是考虑到了当时场景。大家在实际学习工作时,会遇到各种不同实际业务场景,经当时的数据集训练得到的神经网络一般情况下并不能直接应用。
对于一般小型团队或公司而言,项目中能够收集到的训练、测试数据集规模相对较小,相比而言开源社区很多结构复杂的大规模神经网络如ViT使用的数据集则大得多,比如ImageNet系统数据集至少是百万个样本数据。如果直接采用这些”开源神经网络的结构+自己的目标数据集“,很可能会造成过拟合,同时因为此时的目标数据集小太多可能训练得到的最终模型精度(如准确率、泛化能力等)达不到实用程度。
(当然如果某开源模型当时就是由小规模数据集或跟自己当前的目标数据集规模度差不多甚至更小训练得到的,理论上是可能直接拿此开源模型的结构从头开始训练,更新得到目标数据集上的全新权重参数,从而得到一个新的目标模型。)
为此,对于这类大规模神经网络,要想使用它们的结构(甚至进行少量改动),并利用自己当前已有的目标数据集训练、更新权重参数,有如下两种常见的方法:
- 将目标数据集规模、质量扩大或提升至开源的大规模神经网络当时使用的数据集一致。但收集与标数据集是一个成本非常高的过程,其成本大小与数据集规模、质量呈正相关关系,所以这一过程意味更高的成本,这对于一般的小团队或公司是不可取的。
- 使用应用迁移学习(transfer learning),将开源模型在源数据集上学习到的能力迁移到目标数据集上。它的机理是:开源模型使用的源数据集跟目标数据集可能很大程度并不一致或不是描述同类事物,但开源模型在源数据集能够学习到抽象通用的图像特征(比如图像的纹理、边缘、形状与颜色等),它们也适用于目标数据集。
模型方面,应用迁移学习的最常用、最典型应用场景就是模型微调(finetuning)。模型微调的大致意思是:复制某个开源模型(包括模型结构与权重参数),在目标数据集上重新进行训练、不断更新其模型权重参数,使其能够在目标数据集表现达到预期。就是说我们站在巨人的肩膀上,利用优秀开源模型的骨架,赋予其目标数据集的灵魂,使模型与数据完美融合获得预期的涌现能力。
掌握模型微调,可以快速验证问题、提升模型与数据融合能力。
6.3.2 微调在 PyTorch 中的两类典型操作
在实际编写 PyTorch 代码时,根据目标数据集的大小,模型微调通常分为两种具体策略:
- 特征提取(Feature Extraction / 冻结骨干):
- 操作:把开源 Backbone
中除了最后输出层外的所有层参数全冻结(
param.requires_grad = False),只训练我们自己新加的最后一层分类头(Linear)。 - 适用:目标数据集非常非常小(比如只有几百张图)时。
- 操作:把开源 Backbone
中除了最后输出层外的所有层参数全冻结(
具体描述此策略下模型微调的流程如下:
- 1)获取源模型:在源数据集(一般是大规模开源数据集)上训练得到源模型。这个工作一般开源社区已经帮忙完成了。
- 2)构建目标模型:复制源模型中除了最后的输出层外的模型结构与权重参数,并手动构建、添加一个输出与目标数据集一致的新输出层,新输出层的权重参数是随机初始化的。
- 3)更新新输出层:将目标模型中新输出层前面所有的层的权重参数冻结,在目标数据集上进行训练、更新新输出层中的权重参数。
- 4)测试与保存目标模型:当目标模型测试达到预期后,将其以文件的形式持久化保存下来以备后用。如果测试不通过,则重复2)、3)、4)这几个步骤。

全量/部分微调(Full/Partial Fine-tuning):
- 操作:加载预训练权重,用很小的学习率(如 \(10^{-4}\) 或 \(10^{-5}\))对整个网络(或后半部分深层网络)的所有参数进行端到端的梯度更新。
- 适用:目标数据集有一定规模(几千到上万张图)时,也是最标准的微调方式。 在“全量/部分微调”的具体步骤中,通常比“特征提取”多一个学习率差异化(Layer-wise Learning Rate) 的细节:
- 在进行全量微调时,由于 Backbone
已经具备了很好的通用特征,而新替换的输出层(Linear)是随机初始化的:
- 新输出层:可以使用正常的学习率(如 \(10^{-3}\))快速收敛。
- Backbone 预训练层:必须使用极小的学习率(如 \(10^{-5}\) 或 \(10^{-4}\)),避免更新幅度过大直接“破坏”掉原本已经在海量数据上学好的通用特征提取能力(这在学术界被称为“灾难性遗忘”或“破坏预训练特征”)。
在 PyTorch
中,可以通过为优化器传入不同的参数组(param_groups)来实现:
1
2
3
4
5# 差异化学习率配置示例
optimizer = optim.Adam([
{'params': model.backbone.parameters(), 'lr': 1e-5}, # 骨干网络用小学习率
{'params': model.fc.parameters(), 'lr': 1e-3} # 新分类头用正常学习率
])
6.3.4 使用已有模型结构
以下列出了torchvision这一官方开源库中,针对图像识别与分类的常用相关开源模型的实例化:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59"""torchvision 0.21.0 提供的常见图像分类模型结构。
PyTorch 2.6.1 与 torchvision 0.21.0 中,加载预训练模型应使用
weights=某个模型对应的 Weights 枚举值,而不是旧版的 pretrained=True。
本文件会一次创建全部模型实例,便于比较模型结构;但其中 VGG16、
Wide ResNet 等模型占用较多内存。实际实验时通常只保留要使用的一个模型。
"""
import torchvision.models as models # models 集中提供 ResNet、VGG、MobileNet 等视觉模型构造函数。
# ResNet-18:18 层残差网络,通过残差连接缓解深层网络梯度消失;加载 ImageNet 预训练权重。
resnet18 = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# 若不想加载预训练参数,应显式传入 weights=None;这会随机初始化。
# resnet18_without_pretrained_weights = models.resnet18(weights=None)
# AlexNet:经典的早期卷积神经网络,结构简单、参数量较大;此处随机初始化,不加载预训练权重。
alexnet = models.alexnet(weights=None)
# VGG-16:由许多连续的 3×3 卷积构成,结构规整、易理解,但全连接层使参数量和内存占用较大。
vgg16 = models.vgg16(weights=None)
# SqueezeNet 1.0:使用 Fire 模块减少参数量,模型很小,适合内存或部署资源受限的场景。
squeezenet = models.squeezenet1_0(weights=None)
# DenseNet-161:每层接收前面所有层的特征,促进特征复用;网络较深、计算量相对较大。
densenet = models.densenet161(weights=None)
# Inception v3:使用多分支 Inception 模块提取不同尺度特征;训练时默认带辅助分类器。
# init_weights=False 显式选择 PyTorch 的默认层初始化,避免依赖 torchvision 将来会变化的初始化默认值。
inception = models.inception_v3(weights=None, init_weights=False)
# GoogLeNet:较早的 Inception 网络,引入多尺度并行卷积和辅助分类器,参数量相对较少。
# init_weights=False 的作用同上,避免 torchvision 0.21.0 对未指定初始化方式给出 FutureWarning。
googlenet = models.googlenet(weights=None, init_weights=False)
# ShuffleNet V2 x1.0:通过通道拆分和 channel shuffle 降低移动端计算开销,适合轻量化部署。
shufflenet = models.shufflenet_v2_x1_0(weights=None)
# MobileNet V2:使用深度可分离卷积和倒残差结构,以较少计算量获得较好的分类效果。
mobilenet_v2 = models.mobilenet_v2(weights=None)
# MobileNet V3 Large:结合倒残差、SE 注意力和硬激活函数;Large 版本精度优先、模型相对更大。
mobilenet_v3_large = models.mobilenet_v3_large(weights=None)
# MobileNet V3 Small:MobileNet V3 的更小版本,速度和内存占用优先,适合资源更有限的设备。
mobilenet_v3_small = models.mobilenet_v3_small(weights=None)
# ResNeXt-50 32x4d:在残差网络中加入分组卷积(32 个分支),提高网络的“基数”而非只加深或加宽。
resnext50_32x4d = models.resnext50_32x4d(weights=None)
# Wide ResNet-50-2:将 ResNet-50 的通道数加宽约两倍,特征表达能力更强,但参数量和计算量也更高。
wide_resnet50_2 = models.wide_resnet50_2(weights=None)
# MNASNet 1.0:面向移动设备自动搜索得到的轻量网络,在速度与分类精度之间做平衡。
mnasnet = models.mnasnet1_0(weights=None)pretrained=True这种传统老写法,但推荐以下写法:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73"""torchvision 0.21.0 提供的常见图像分类模型结构。
PyTorch 2.6.1 与 torchvision 0.21.0 中,加载预训练模型应使用
weights=某个模型对应的 Weights 枚举值,而不是旧版的 pretrained=True。
本文件会一次创建全部模型实例,便于比较模型结构;但其中 VGG16、
Wide ResNet 等模型占用较多内存。实际实验时通常只保留要使用的一个模型。
"""
import torchvision.models as models # models 集中提供 ResNet、VGG、MobileNet 等视觉模型构造函数。
# ResNet-18:18 层残差网络,通过残差连接缓解深层网络梯度消失;加载 ImageNet 预训练权重。
resnet18 = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# 若不想加载预训练参数,应显式传入 weights=None;这会随机初始化。
# resnet18_without_pretrained_weights = models.resnet18(weights=None)
# AlexNet:经典的早期卷积神经网络,结构简单、参数量较大;加载 ImageNet 预训练权重。
alexnet = models.alexnet(weights=models.AlexNet_Weights.DEFAULT)
# VGG-16:由许多连续的 3×3 卷积构成,结构规整、易理解,但全连接层使参数量和内存占用较大;加载 ImageNet 预训练权重。
vgg16 = models.vgg16(weights=models.VGG16_Weights.DEFAULT)
# SqueezeNet 1.0:使用 Fire 模块减少参数量,模型很小,适合内存或部署资源受限的场景;加载 ImageNet 预训练权重。
squeezenet = models.squeezenet1_0(weights=models.SqueezeNet1_0_Weights.DEFAULT)
# DenseNet-161:每层接收前面所有层的特征,促进特征复用;网络较深、计算量相对较大;加载 ImageNet 预训练权重。
densenet = models.densenet161(weights=models.DenseNet161_Weights.DEFAULT)
# Inception v3:使用多分支 Inception 模块提取不同尺度特征;训练时默认带辅助分类器;加载 ImageNet 预训练权重。
# init_weights=False 显式选择 PyTorch 的默认层初始化,避免依赖 torchvision 将来会变化的初始化默认值。
inception = models.inception_v3(
weights=models.Inception_V3_Weights.DEFAULT, init_weights=False
)
# GoogLeNet:较早的 Inception 网络,引入多尺度并行卷积和辅助分类器,参数量相对较少;加载 ImageNet 预训练权重。
# init_weights=False 的作用同上,避免 torchvision 0.21.0 对未指定初始化方式给出 FutureWarning。
googlenet = models.googlenet(
weights=models.GoogLeNet_Weights.DEFAULT, init_weights=False
)
# ShuffleNet V2 x1.0:通过通道拆分和 channel shuffle 降低移动端计算开销,适合轻量化部署;加载 ImageNet 预训练权重。
shufflenet = models.shufflenet_v2_x1_0(
weights=models.ShuffleNet_V2_X1_0_Weights.DEFAULT
)
# MobileNet V2:使用深度可分离卷积和倒残差结构,以较少计算量获得较好的分类效果;加载 ImageNet 预训练权重。
mobilenet_v2 = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT)
# MobileNet V3 Large:结合倒残差、SE 注意力和硬激活函数;Large 版本精度优先、模型相对更大;加载 ImageNet 预训练权重。
mobilenet_v3_large = models.mobilenet_v3_large(
weights=models.MobileNet_V3_Large_Weights.DEFAULT
)
# MobileNet V3 Small:MobileNet V3 的更小版本,速度和内存占用优先,适合资源更有限的设备;加载 ImageNet 预训练权重。
mobilenet_v3_small = models.mobilenet_v3_small(
weights=models.MobileNet_V3_Small_Weights.DEFAULT
)
# ResNeXt-50 32x4d:在残差网络中加入分组卷积(32 个分支),提高网络的“基数”而非只加深或加宽;加载 ImageNet 预训练权重。
resnext50_32x4d = models.resnext50_32x4d(
weights=models.ResNeXt50_32X4D_Weights.DEFAULT
)
# Wide ResNet-50-2:将 ResNet-50 的通道数加宽约两倍,特征表达能力更强,但参数量和计算量也更高;加载 ImageNet 预训练权重。
wide_resnet50_2 = models.wide_resnet50_2(
weights=models.Wide_ResNet50_2_Weights.DEFAULT
)
# MNASNet 1.0:面向移动设备自动搜索得到的轻量网络,在速度与分类精度之间做平衡;加载 ImageNet 预训练权重。
mnasnet = models.mnasnet1_0(weights=models.MNASNet1_0_Weights.DEFAULT)
相关注意事项:
- 对于PyTorch内置的每个模型,其官方都维护了一个存储与下载地址并在相关源码中记录着,一般以
.pth格式存在。当尝试实例化模型时,如果本地没有找到(会进行md5值校验)此.pth文件源码会自动下载,找到了就直接使用。 - 除了从PyTorch提供的官方地址下载模型文件,PyTorch还提供了从自定义地址下载模型文件、保存到自定义本地位置(默认下载的模型文件保存在
用户家目录/.cache/torch/hub/checkpoint目录下)等,如使用torch.utils.model_zoo.load_url方法,其官方文档介绍如:torch.utils.model_zoo.load_url 。用法示例如下所示。
1 | |
6.3.5 训练特定层
默认情况下,模型中具体某个模型块或层的权重或偏置的requires_grad属性值为True,表示其在损失反向传播的过程中会记录梯度、被优化器自动更新,这在模型从头开始训练或全量微调过程中是正常无需修改的。
但正如上面所述,如果我们基于某个大规模开源模型构造了目标模型进行微调,其最后输出层是重新构建的(新构建的输出层默认权重与偏置的requires_grad属性值默认为True)、前面的层结构和权重参数是继承自原开源模型。此时就需要将目标模型最后输出层前面的层的权重参数进行冻结,以保证在微调过程中只更新输出层的权重参数,此时就需要将前面的层的权重或偏置的requires_grad属性值强制更改为False,表示在微调中不更新它们。
1 | |
此示例代码基于resnet18开源模型为基础构建目标模型,将目标模型的最后一层由原来的输出1000个分类更改为输出4个分类,并将输出层前面层的权重或偏置的requires_grad属性值强制更改为False以冻结它们。但未完整给出后续微调具体代码。
6.4 模型微调 - timm
通过前面的学习,已知PyTorch官方的torchvision库中包含了许多开源模型。除此之外,还有一个开源timm中也包含了大量开源模型、与工具和优化器等功能。 timm(PyTorch Image Models)是一组图像模型、层、实用工具、优化器、调度器、数据加载器/增强方法,以及参考的训练和验证脚本,旨在整合多种最先进的模型,使其能够重现 ImageNet 的训练结果。
timm开源项目中包含的开源模型是torchvision中开源模型的父集。
它的github官方链接是:https://github.com/huggingface/pytorch-image-models(还在更新中),官方文档网站是:https://huggingface.co/docs/hub/timm、https://timm.fast.ai/(此是老网站,最近4年无更新)。
除了官方文档,还有一篇个人博客整理了timm中常见的用法:https://www.aiuai.cn/aifarm1967.html,可以参考。
6.4.1 timm的安装
根据官方文档,其提供了如下两种安装方式:
(1)使用pip安装 1
pip install timm1
2git clone https://github.com/huggingface/pytorch-image-models.git
cd pytorch-image-models && pip install -e .
6.4.2 timm的使用
timm/__init__.py
将常用模型注册表接口重新导出,因此通常只需import timm即可使用它们。
下面按照“查找模型 → 确认模型和权重 → 查看配置 →
创建模型”的顺序介绍常用方法。
模型和权重配置会随安装的timm版本变化,示例中的名称应以本机查询结果为准。
1 | |
6.4.2.1 查询可用模块与模型
查询可用模块与模型,分别使用list_modules()和list_models()
方法。
timm.list_modules()返回已注册模型所在的模块名列表,返回值是按字母排序的list[str]。
它不是Python模块对象或文件路径,通常用于为list_models()的module参数选择范围。
1 | |
timm.list_models()只读取模型注册表,不会创建模型,也不会下载权重。
其返回值是按自然顺序排序的模型名称列表。
1 | |
可以同时按模块和名称筛选;两项条件需要同时满足。多个module也是并集。
1 | |
pretrained=True表示只保留具有有效预训练权重配置的模型。此时include_tags默认自动设为True,
所以结果会带有具体权重标签,例如resnet50.a1_in1k:前半部分是架构名,点号后的部分是权重配置标签。
列出模型不会下载这些权重。
1 | |
timm.list_pretrained()是查询具体预训练权重名称的快捷方法,等价于
list_models(pretrained=True, include_tags=True, ...)。当需要排除权重名称时,建议把
exclude_filters写成列表,即使只有一个规则。
1 | |
若批量测试模型而不希望同一架构的兼容别名重复出现,可加入name_matches_cfg=True。
日常选择模型时通常不需要该参数。
1 | |
6.4.2.2 确认模型与预训练权重
确认模型与预训练权重分别使用方法:is_model()和is_model_pretrained()。
在把字符串形式的模型名称交给训练脚本前,可以先进行检查。is_model()检查架构是否已注册;
即使传入带权重标签的名称,它也会识别其中的架构部分。is_model_pretrained()则检查该名称
是否有有效的预训练权重配置。
1 | |
这两个检查只查询本地注册表;它们不能保证权重文件已经缓存到本地,也不能保证网络一定能够 下载权重。
6.4.2.3 查看预训练配置
查看预训练配置主要使用方法:get_pretrained_cfg()和get_pretrained_cfg_value()。
一套预训练权重除了参数文件外,还对应输入尺寸、归一化均值和标准差、裁剪比例、类别数等配置。
get_pretrained_cfg()返回一个PretrainedCfg对象;传入带标签的名称可以精确查询某一套权重。
它只返回配置副本,不下载权重。
1 | |
如果只需要一个字段,可使用get_pretrained_cfg_value()。字段不存在时它返回None,
适合在程序中提供兼容处理。
1 | |
应使用list_pretrained()或list_models(pretrained=True)返回的名称来查询配置。
对于不存在的架构或错误的权重标签,配置查询会报错;这能及早发现名称拼写问题。
6.4.2.4 创建与修改模型
创建模型使用的方法是create_model() 。
timm.create_model()根据模型名称查找构造函数,并返回一个torch.nn.Module。
pretrained=False(默认值)只创建随机初始化的模型;pretrained=True才会加载该模型或指定标签
对应的预训练权重。第一次加载预训练权重时,可能需要从网络下载并缓存文件。
1 | |
迁移学习时,最常见的是指定预训练模型并修改分类类别数。下面的num_classes=4会创建一个
输出4类的新分类头;预训练的主干网络参数会被加载,而新分类头需要在自己的数据集上训练。
1 | |
若只希望按架构默认的预训练配置加载,也可以使用不带标签的名称:
1 | |
甚至timm在创建模型实例时可以改变其输入通道(实现原理参考:https://fastai.github.io/timmdocs/models#So-how-is-timm-able-to-load-these-weights?):
1
2
3
4
5
6
7
8
9
10model = timm.create_model(
'resnet50d',
pretrained=True,
in_chans=1 # resnet50d 构架的模型原本的输入通道是3,此处改为1
)
print(model)
# 构造一个单通道测试数据(输入数据的空间尺寸跟原开源模型要求的一致)
x = troch.randn(1, 1, 224, 224)
out = model(x) print(out.shape) # 结果形状: torch.Size([1, 1000])
6.4.2.5 获取模型构造函数
获取模型构造函数使用方法:model_entrypoint() 。
timm.model_entrypoint()返回已注册模型对应的构造函数(callable),而不是模型实例。
它主要用于阅读源码、实现模型工厂或编写框架扩展;普通训练代码优先使用create_model(),
因为后者还会处理权重标签、外部检查点和部分通用配置。
1 | |
综合来看,日常任务通常使用list_models()或list_pretrained()发现名称,使用is_model()确认名称,再通过create_model()创建模型,当需要复现预处理或核对某套权重的训练设定时,
再查询get_pretrained_cfg()。
6.4.4 模型的保存
利用timm创建的模型实例的类型是timm.models.xx.yy,但这个类型又是继承自torch.nn.Module。
所以说利用timm创建的模型实例的类型本质上仍是PyTorch中torch.nn.Module,那么就可以使用PyTorch中提供的方法的来保存模型。
1 | |
6.5 半精度与混合精度训练
6.5.1 半精度与混合精度训练的背景
GPU 的性能主要由算力(Flops,在特定精度下的矩阵乘加运算吞吐量)与显存(包括容量大小与显存带宽)共同决定。对于一张确定的物理 GPU,其峰值算力与显存硬件规格是固定的。
显存作为 GPU 板载的高速存储介质,用于在计算过程中高速暂存模型参数、输入数据及中间激活结果。 在GPU算力范围内,GPU单次能接受的输入数据量越多(常以增加batch_size值大小体现),通常意味着单个轮次模型训练所需要的迭代步数越少、模型训练总耗时越少。GPU内显存中默认是以float32单精度浮点数暂存各种数据的,模型训练的一个轮次中,在float32单精度情况下当batch_size增大一定程度时,输入数据需要的显存可能达到GPU显存的临时存储上限,如果此时所需的GPU算力还未达到GPU的算力上限,为了进一步压榨GPU的计算能力与显存临时存储能力、提升模型的训练效率,可以float16半精度浮点数的形式将数据送入GPU显存中,此时同样的源数据只需要占据使用float32时的大约一半GPU显存空间,所以能使得batch_size最大增加为原来的2倍左右、提升模型的训练效率。
但因为float16相比float32使用了更少的比特位来存储数据,指数位数与小数位数都更少,所以能表达的数据范围也更小。但对于一般的模型训练场景下用float16来存储数据已经够用了,所以对于一些单样本占据数据量比较大(如3D图像、视频等)、数据集总样本数很多的场景,使用半精度进行模型训练仍然一种常用的提升模型训练效率的方法。
关于float32与float16的相关示意图如下:
可以看到:
float32单精度使用4个字节存储一个数据,其中符号位、指数位、小数位使用的比特位数分别是:1、8、23 。存储的数据范围是:\(-3.40 \times 10^{38} \sim 3.40 \times 10^{38}\) 。
float16半精度使用2个字节存储一个数据,其中符号位、指数位、小数位使用的比特位数分别是:1、5、10 。存储的数据范围是:\(-65504 \sim 65504\) 。
如果float16表达的数据范围不够用的,但又不想使用float32,也可以使用BF16,它能够表达的数据范围与float32几乎一致,只是减少了尾数位、牺牲精度。
此外,在实际的大模型或深度学习训练中,如果单纯直接将所有数据和参数都转换为 FP16 进行训练,很容易遇到数值不稳定(如梯度消失或梯度溢出)的问题。为了解决这一问题,业界普遍采用 AMP(Automatic Mixed Precision,自动混合精度) 策略——在前向/反向传播中采用 FP16/BF16 进行矩阵计算以提升效率,同时在后台维护一份 FP32 的主权重并结合动态梯度缩放(Loss Scaling)来保证数值稳定性,兼顾了训练效率与模型精度。
6.5.2 半精度训练示例
在PyTorch中使用半精度进行训练,只需要在使用默认float32单精度的基础修改以下几处即可:
- 导入autocast:from torch.cuda.amp import autocast 。 -
用装饰器@autocast()装饰模型类的forward方法 。 -
模型训练时,将数据前向传播及损失计算代码放在with autocast():语句块中。
以下是一个完整可执行的使用半精度进行模型训练的示例:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171"""使用旧版 torch.cuda.amp.autocast 接口进行“半精度训练”的示例。
本文件刻意对应旧教程的写法:
1. 从 torch.cuda.amp 导入 autocast;
2. 使用 @autocast() 装饰模型的 forward();
3. 在训练阶段使用 with autocast()。
注意:autocast 的准确名称是“自动混合精度(AMP)”,并非让所有运算都变为
float16。PyTorch 2.6 已推荐 torch.autocast 与 torch.amp.GradScaler;这里保留
旧接口仅用于学习和对应教程。若没有 CUDA GPU,旧接口会自动禁用,代码改为 FP32
训练,仍可直接运行。
"""
import torch
from torch import nn
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader, TensorDataset
# 固定随机种子,使数据、模型初始参数和训练结果基本可复现。
torch.manual_seed(42)
# CUDA 可用时将模型和数据移动到 GPU;否则在 CPU 上运行。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 旧版 autocast 和 GradScaler 仅在 CUDA 可用时启用;CPU 上仍保留相同代码结构,
# 但会透明地执行普通 FP32 训练。
legacy_amp_enabled = torch.cuda.is_available()
def make_dataset():
"""生成不需要下载的二维二分类数据集。"""
# features 的形状为 (400, 2),每一行是一个包含两个 float32 特征的样本。
features = torch.randn(400, 2)
# 通过一条直线生成类别标签;标签的形状为 (400,),取值为 0 或 1。
labels = (1.5 * features[:, 0] - features[:, 1] > 0).long()
# 翻转 20 个标签,模拟真实任务中少量标注噪声,避免任务过于理想化。
noisy_indices = torch.randperm(features.size(0))[:20]
labels[noisy_indices] = 1 - labels[noisy_indices]
# 前 320 条用于训练,后 80 条用于验证。
return (
TensorDataset(features[:320], labels[:320]),
TensorDataset(features[320:], labels[320:]),
)
class Classifier(nn.Module):
"""输入两个特征、输出两个类别 logits 的小型全连接分类网络。"""
def __init__(self):
super().__init__()
self.network = nn.Sequential(
nn.Linear(2, 16), # 将 2 个输入特征映射为 16 个隐藏特征。
nn.ReLU(), # 引入非线性,使模型能学习更复杂的决策边界。
nn.Linear(16, 2), # 输出每个类别的原始分数,形状为 (batch_size, 2)。
)
# 这是旧教程常见的装饰器写法:执行 forward 时自动进入 CUDA autocast 上下文。
# 在 PyTorch 2.6 中它会产生弃用警告;第 51 个文件会展示推荐的新写法。
@autocast(enabled=legacy_amp_enabled)
def forward(self, features):
"""执行前向传播,返回尚未经过 softmax 的 logits。"""
return self.network(features)
def evaluate(model, data_loader):
"""在验证集上计算正确率,不更新模型参数。"""
model.eval()
correct = 0
total = 0
# 验证阶段不需要梯度,因此关闭梯度记录。
with torch.no_grad():
for features, labels in data_loader:
features = features.to(device)
labels = labels.to(device)
# 旧教程也会在前向计算外显式使用 with autocast()。
# 即使 forward 已有装饰器,这样写仍然可行;嵌套上下文不会改变训练逻辑。
with autocast(enabled=legacy_amp_enabled):
logits = model(features)
predictions = logits.argmax(dim=1)
correct += (predictions == labels).sum().item()
total += labels.size(0)
return correct / total
def main():
train_dataset, valid_dataset = make_dataset()
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
model = Classifier().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.03)
# 旧版 GradScaler 与旧版 autocast 配套使用:它在 CUDA FP16 时放大 loss,
# 防止过小梯度在 float16 中下溢为 0;无 CUDA 时会自动禁用。
scaler = GradScaler(enabled=legacy_amp_enabled)
print(f"运行设备: {device}")
print(f"CUDA 是否可用:{legacy_amp_enabled},旧版 autocast 是否会启用: {legacy_amp_enabled}")
epochs = 30
for epoch in range(epochs):
model.train()
total_loss = 0.0
total_samples = 0
for features, labels in train_loader:
features = features.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# 旧教程要求将前向传播与损失计算放在 with autocast() 中。
# CUDA 上适合的算子会使用 float16;数值敏感算子仍可能保持 float32。
with autocast(enabled=legacy_amp_enabled):
logits = model(features)
loss = loss_fn(logits, labels)
# 使用旧版 GradScaler 完成:放大 loss -> 反向传播 -> 更新参数 -> 更新缩放比例。
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
batch_size = labels.size(0)
total_loss += loss.item() * batch_size
total_samples += batch_size
train_loss = total_loss / total_samples
valid_accuracy = evaluate(model, valid_loader)
if (epoch + 1) % 5 == 0 or epoch == 0:
print(
f"Epoch [{epoch + 1:2d}/{epochs}] "
f"train_loss: {train_loss:.4f} "
f"valid_accuracy: {valid_accuracy:.2%}"
)
if __name__ == "__main__":
main()
"""
### 带 NVIDIA 的服务器上某次执行:
(py31213) root@controller01:~# python 50-半精度训练.py
/opt/installPkgs/research-PyTorch/50-半精度训练.py:63: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
@autocast(enabled=legacy_amp_enabled)
/opt/installPkgs/research-PyTorch/50-半精度训练.py:104: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.
scaler = GradScaler(enabled=legacy_amp_enabled)
运行设备: cuda
CUDA 是否可用:True,旧版 autocast 是否会启用: True
/opt/installPkgs/research-PyTorch/50-半精度训练.py:122: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
with autocast(enabled=legacy_amp_enabled):
/opt/installPkgs/research-PyTorch/50-半精度训练.py:83: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
with autocast(enabled=legacy_amp_enabled):
Epoch [ 1/30] train_loss: 0.4509 valid_accuracy: 85.00%
Epoch [ 5/30] train_loss: 0.2467 valid_accuracy: 86.25%
Epoch [10/30] train_loss: 0.2394 valid_accuracy: 83.75%
Epoch [15/30] train_loss: 0.2356 valid_accuracy: 87.50%
Epoch [20/30] train_loss: 0.2367 valid_accuracy: 87.50%
Epoch [25/30] train_loss: 0.2264 valid_accuracy: 87.50%
Epoch [30/30] train_loss: 0.2204 valid_accuracy: 85.00%
"""
半精度训练主要适用于数据本身的size比较大(比如说3D图像、视频等)。当数据本身的size并不大时(比如手写数字MNIST数据集的图片尺寸只有28*28),使用半精度训练则可能不会带来显著的提升。
其他示例,可参考:https://featurize.cn/notebooks/368cbc81-2b27-4036-98a1-d77589b1f0c4。
6.5.3 混合精度训练及示例
混合精度训练(Automatic Mixed Precision,简称 AMP) 是现代深度学习训练(尤其是大语言模型 LLM 和 Vision Transformer)中用于大幅提升训练速度并显著降低显存占用的标准工程实践。
深度学习中的现代混合精度训练(Mixed Precision
Training)方法是由百度(Baidu)与英伟达(NVIDIA)的研究团队于2017年底联合提出,论文见:https://arxiv.org/abs/1710.03740。
它的核心思想非常简单:在对精度不敏感的大规模计算环节使用 低精度数据类型(FP16/BF16) 来换取算力和显存效率,而在对数值范围敏感的参数更新环节保持 高精度数据类型(FP32) 来确保模型的收敛与稳定性。
6.5.3.1 为什么需要混合精度?
- FP32(单精度浮点数)的局限:
- 显存大、算力未跑满:FP32 占用 4 字节。当模型参数量达到数亿或数十亿时,FP32 会迅速榨干 GPU 显存。
- 纯 FP16(半精度浮点数)的致命缺陷:
- 数值范围过窄(\(-65504 \sim 65504\))。FP16 的指数位只有 5 位,最小正数约为 \(6.10 \times 10^{-5}\)。
- 梯度下溢(Underflow)与上溢(Overflow):反向传播计算出的微小梯度低于
\(10^{-5}\) 时会被直接清零(下溢,变为
0);如果数值大于 65504 则会溢出(变为
Inf/NaN),导致模型无法训练。
为了兼顾 FP16 的极速/省显存 与 FP32 的稳定性,混合精度训练(AMP)应运而生。
6.5.3.2 混合精度训练(AMP)的核心运行机制
AMP 通常由以下三个核心技术环节协同工作。以 PyTorch 原生 AMP
为例,模型参数默认仍保存为 FP32;autocast
不会把整个模型永久转换为 FP16/BF16,而是按算子临时选择合适的计算精度:
1
2
3
4
5
6
7
8
9 前向传播(autocast 按算子选择精度)
FP32 输入、FP32 模型参数 ──> 兼容的矩阵乘/卷积:FP16 或 BF16
数值敏感的算子:FP32
│
▼
loss 与反向传播计算梯度
│
▼
FP32 模型参数与通常为 FP32 的优化器状态更新
1. 前向与反向传播(按算子选择低精度)
- 计算载体:输入和模型参数无需预先整体转换为
FP16/BF16。进入
autocast上下文后,PyTorch 会为每个算子选择精度:卷积、矩阵乘法等兼容算子通常以 FP16 或 BF16 计算;归约等数值敏感算子则保持 FP32。中间激活值的实际 dtype 也因此取决于产生它的算子。 - 硬件加速:在支持 Tensor Core 的 CUDA GPU 上,低精度矩阵乘法和卷积通常能提升吞吐量,并降低部分激活值的显存占用;实际加速比例取决于 GPU、模型结构、批量大小和数据输入瓶颈,不能保证固定倍数。
2. FP32 模型参数与优化器状态
- 更新载体:在 PyTorch 原生 AMP
的常见用法中,模型通过
model.to(device)保持 FP32 参数,Adam/AdamW 等优化器的状态也通常随 FP32 参数以 FP32 保存。 - 作用:优化器直接用这些 FP32
参数进行更新,避免把参数永久存为 FP16/BF16
后的频繁舍入误差。某些其他框架或手动
model.half()的实现会维护独立的“FP32 主权重”副本,但这不是torch.autocast的默认工作方式。
3. 动态梯度缩放(Dynamic Loss Scaling,针对 FP16)
- 机制:在反向传播前,将损失函数(Loss)乘以一个较大的放大因子 \(S\)(Scale Factor),强制将极小的梯度拉回到 FP16 的表达范围内,避免下溢。
- 还原:
GradScaler.step(optimizer)会在调用优化器前还原梯度,并检查是否出现Inf/NaN。若发现溢出,则跳过本次参数更新,并由GradScaler.update()动态调小 \(S\)。
6.5.3.3 FP16 AMP 与 BF16 AMP 的选型对比
| 存储格式 | 字节数 | 指数位 / 尾数位 | 表达数据范围 | 稳定性与特点 | 适用场景 |
|---|---|---|---|---|---|
| FP16 AMP | 2 字节 | \(5\text{ bit} / 10\text{ bit}\) | \(-65504 \sim 65504\) | 精度略高,但范围小;必须开启
GradScaler(损失缩放)。 |
传统 CV 视觉模型、较早期的 GPU(如 V100/T4/20系) |
| BF16 AMP | 2 字节 | \(8\text{ bit}\) \(/ 7\text{ bit}\) | \(\approx \pm 3.4 \times 10^{38}\) | 范围与 FP32 完全一致;极难溢出,无需
GradScaler。 |
大语言模型(LLM)、ViT、主流现代 GPU(RTX 30/40系、A100/H100) |
6.5.3.4 混合精度训练的工程优势
- 显存消耗显著降低(减少约 40% ~ 50%):
- 占显存大头的数据(激活值、中间张量)全部变为 16 位,允许提升 1.5~2
倍的
batch_size。
- 占显存大头的数据(激活值、中间张量)全部变为 16 位,允许提升 1.5~2
倍的
- 训练计算速度大幅提升(通常 1.5x ~ 3x 加速):
- 极大地缓解了 GPU 显存带宽瓶颈(Memory Bandwidth Bound),同时全面释放了 Tensor Core 的矩阵计算吞吐量。
- 模型最终精度几乎无损:
- 由于关键的权重更新与损失计算依然运行在 FP32 精度下,模型的收敛曲线与最终准确率/Loss 与全 FP32 训练高度吻合。
6.5.3.5 PyTorch 代码实现范式
在 PyTorch 2.6 中,推荐使用 torch.autocast 与
torch.amp.GradScaler。下面只保留训练循环中与 CUDA FP16 AMP
直接相关的关键代码;无 CUDA 时通过 enabled=False
自动退化为普通 FP32 训练:
1 | |
以下是完整的可执行示例1(使用了torch.float16):
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181"""使用 PyTorch 2.6 推荐接口进行 AMP(自动混合精度)训练的完整示例。
与第 50 个文件的旧教程写法不同,本例使用:
- torch.autocast(device_type="cuda", dtype=torch.float16)
- torch.amp.GradScaler("cuda")
在 CUDA GPU 上,autocast 会让适合低精度的算子使用 float16,同时让模型参数和
数值敏感算子保持 float32;GradScaler 用于避免小梯度在 float16 中下溢为 0。
没有 CUDA GPU 时,amp_enabled=False,代码自动退化为普通 FP32 训练。
"""
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
# 固定随机种子,使数据、模型初始参数和训练结果基本可复现。
torch.manual_seed(42)
# CUDA GPU 可用时使用它;否则使用 CPU。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 本例实现的是 CUDA FP16 AMP,因此仅在 CUDA 环境启用。
amp_enabled = device.type == "cuda"
def make_dataset():
"""生成不需要下载的二维二分类数据集。"""
# features 的形状为 (400, 2):400 个样本,每个样本有两个 float32 特征。
features = torch.randn(400, 2)
# 根据线性边界生成 0/1 标签,并加入少量噪声使任务更接近真实情形。
labels = (1.5 * features[:, 0] - features[:, 1] > 0).long()
# torch.randperm(N): 生成一个包含从 0 到 N-1 的所有整数的一维张量(Tensor),并将这些整数的顺序随机打乱。
# torch.randperm(N)[:20]: 从包含 0 到 N-1 的所有整数中,随机无放回地抽取 20 个不重复的整数。
noisy_indices = torch.randperm(features.size(0))[:20]
labels[noisy_indices] = 1 - labels[noisy_indices]
# 前 320 条为训练集,后 80 条为验证集。
return (
TensorDataset(features[:320], labels[:320]),
TensorDataset(features[320:], labels[320:]),
)
class Classifier(nn.Module):
"""输入两个特征、输出两个类别 logits 的小型全连接分类网络。"""
def __init__(self):
super().__init__()
self.network = nn.Sequential(
nn.Linear(2, 16), # 输入维度 2,隐藏层维度 16。
nn.ReLU(), # 增加非线性表达能力。
nn.Linear(16, 2), # 输出两个类别的原始分数,形状为 (batch_size, 2)。
)
def forward(self, features):
# 不使用 @autocast 装饰器:在训练循环中控制 autocast,
# 可以分别决定训练、验证和推理阶段是否开启混合精度。
return self.network(features)
def evaluate(model, data_loader):
"""在验证集上计算分类正确率。"""
model.eval()
correct = 0
total = 0
# 验证不需要计算梯度,关闭梯度记录可节省内存。
with torch.no_grad():
for features, labels in data_loader:
features = features.to(device)
labels = labels.to(device)
# 验证阶段也可以使用 autocast;无 CUDA 时 enabled=False,实际仍为 FP32。
with torch.autocast(
device_type="cuda", dtype=torch.float16, enabled=amp_enabled
):
logits = model(features)
predictions = logits.argmax(dim=1)
correct += (predictions == labels).sum().item()
total += labels.size(0)
return correct / total
def main():
train_dataset, valid_dataset = make_dataset()
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
# AMP 不要求将模型整体转为 float16;参数默认保留 float32。
model = Classifier().to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.03)
# GradScaler 仅在 CUDA FP16 AMP 启用时实际缩放梯度。
# 当 amp_enabled=False(本机没有 CUDA GPU)时,enabled=False:
# scaler 不会放大 loss,也不会检查或跳过参数更新,而是作为“直通包装器”使用,
# 因此下面的 scale().backward()、step()、update() 与普通 FP32 训练等价。
scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
print(f"运行设备: {device}")
print(f"是否启用 CUDA float16 AMP: {amp_enabled}")
epochs = 30
for epoch in range(epochs):
model.train()
total_loss = 0.0
total_samples = 0
for features, labels in train_loader:
features = features.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# autocast 仅包裹前向传播和损失计算,反向传播与 optimizer.step() 放在其外部。
#
# amp_enabled 是总开关:
# - 纯 CPU 环境中,device.type == "cpu",所以 amp_enabled=False。
# 此时 enabled=False 会禁用 autocast;这个 with 语句只是一个“不做任何
# dtype 转换”的上下文,features、模型参数、logits、loss 和梯度都保持
# 默认的 float32,训练过程就是普通 FP32 训练。
# - CUDA 环境中,amp_enabled=True。autocast 不会先把 features 或整个 model
# 永久转换为 float16;它会在执行每个算子时自行决定计算 dtype。通常卷积、
# 矩阵乘法等适合 Tensor Core 的算子会使用 float16,而归约等数值敏感算子会
# 保持 float32。因此,某个中间张量的 dtype 取决于产生它的算子。
with torch.autocast(
device_type="cuda", dtype=torch.float16, enabled=amp_enabled
):
logits = model(features) # 在 CUDA AMP 时,模型内部的兼容算子会自动使用 float16。
loss = loss_fn(logits, labels) # 损失函数内部的数值敏感部分会按 autocast 规则保持合适精度。
# 下面三步主要解决 CUDA FP16 的“小梯度下溢”问题:
# float16 能表示的数值范围较小,极小梯度可能变成 0。scaler 会先将 loss
# 乘以一个动态缩放因子,使梯度也按同一比例变大,从而更可能落在 float16
# 的可表示范围内;这不会改变真正要更新的梯度方向。
scaler.scale(loss).backward()
# 在调用 optimizer.step() 前,scaler 会先将参数梯度除以缩放因子,恢复真实
# 大小,然后检查是否有 Inf 或 NaN。若发现溢出,会跳过本批次的参数更新,
# 防止错误梯度破坏模型参数;若正常,则调用 optimizer 更新默认的 FP32 参数。
scaler.step(optimizer)
# 根据刚才是否发生溢出,动态调整下一批次使用的缩放因子:
# 发生溢出时通常缩小缩放因子;连续若干批正常时可逐步增大。
scaler.update()
batch_size = labels.size(0)
total_loss += loss.item() * batch_size
total_samples += batch_size
train_loss = total_loss / total_samples
valid_accuracy = evaluate(model, valid_loader)
if (epoch + 1) % 5 == 0 or epoch == 0:
print(
f"Epoch [{epoch + 1:2d}/{epochs}] "
f"train_loss: {train_loss:.4f} "
f"valid_accuracy: {valid_accuracy:.2%}"
)
if __name__ == "__main__":
main()
"""
### 带 NVIDIA 的服务器上某次执行:
(py31213) root@controller01:~# python 51-amp混合精度训练-fp16.py
运行设备: cuda
是否启用 CUDA float16 AMP: True
Epoch [ 1/30] train_loss: 0.4509 valid_accuracy: 85.00%
Epoch [ 5/30] train_loss: 0.2467 valid_accuracy: 86.25%
Epoch [10/30] train_loss: 0.2394 valid_accuracy: 83.75%
Epoch [15/30] train_loss: 0.2356 valid_accuracy: 87.50%
Epoch [20/30] train_loss: 0.2367 valid_accuracy: 87.50%
Epoch [25/30] train_loss: 0.2264 valid_accuracy: 87.50%
Epoch [30/30] train_loss: 0.2204 valid_accuracy: 85.00%
"""
以下是完整的可执行示例2(使用了torch.bfloat16):
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184"""使用 PyTorch 2.6 推荐接口进行 CUDA BF16 AMP 训练的完整示例。
本例与 51-amp混合精度训练-fp16.py 的区别是低精度 dtype:
- FP16 示例使用 torch.float16,通常需要 GradScaler 防止小梯度下溢;
- 本例使用 torch.bfloat16(BF16),其指数范围接近 float32,通常不需要 GradScaler。
在支持 BF16 的 CUDA GPU 上,autocast 会让兼容算子使用 BF16,同时让模型参数和
数值敏感算子保持 float32。没有 CUDA GPU,或 CUDA GPU 不支持 BF16 时,
amp_enabled=False,代码自动退化为普通 FP32 训练。
"""
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
# 固定随机种子,使数据、模型初始参数和训练结果基本可复现。
torch.manual_seed(42)
# CUDA GPU 可用时使用它;否则使用 CPU。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# BF16 AMP 既需要 CUDA,也需要 GPU 本身支持 BF16 指令。
# “and” 使用短路求值:没有 CUDA 时不会执行 torch.cuda.is_bf16_supported()。
bf16_supported = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
# 这是 autocast 的总开关。False 时,下面的 autocast 上下文不会转换任何 dtype。
amp_enabled = bf16_supported
# 明确指定本例的低精度类型。BF16 与 float32 有同样的 8 位指数位,
# 数值范围接近 float32,但尾数位更少,因此精度低于 float32。
amp_dtype = torch.bfloat16
def make_dataset():
"""生成不需要下载的二维二分类数据集。"""
# features 的形状为 (400, 2):400 个样本,每个样本有两个 float32 特征。
features = torch.randn(400, 2)
# 根据线性边界生成 0/1 标签,并加入少量噪声使任务更接近真实情形。
labels = (1.5 * features[:, 0] - features[:, 1] > 0).long()
# torch.randperm(N): 生成一个包含从 0 到 N-1 的所有整数的一维张量(Tensor),并将这些整数的顺序随机打乱。
# torch.randperm(N)[:20]: 从包含 0 到 N-1 的所有整数中,随机无放回地抽取 20 个不重复的整数。
noisy_indices = torch.randperm(features.size(0))[:20]
labels[noisy_indices] = 1 - labels[noisy_indices]
# 前 320 条为训练集,后 80 条为验证集。
return (
TensorDataset(features[:320], labels[:320]),
TensorDataset(features[320:], labels[320:]),
)
class Classifier(nn.Module):
"""输入两个特征、输出两个类别 logits 的小型全连接分类网络。"""
def __init__(self):
super().__init__()
self.network = nn.Sequential(
nn.Linear(2, 16), # 输入维度 2,隐藏层维度 16。
nn.ReLU(), # 增加非线性表达能力。
nn.Linear(16, 2), # 输出两个类别的原始分数,形状为 (batch_size, 2)。
)
def forward(self, features):
# 不使用 @autocast 装饰器:在训练循环中控制 autocast,
# 可以分别决定训练、验证和推理阶段是否开启混合精度。
return self.network(features)
def evaluate(model, data_loader):
"""在验证集上计算分类正确率。"""
model.eval()
correct = 0
total = 0
# 验证不需要计算梯度,关闭梯度记录可节省内存。
with torch.no_grad():
for features, labels in data_loader:
features = features.to(device)
labels = labels.to(device)
# 验证阶段也可以使用 BF16 autocast;无 CUDA 或不支持 BF16 时,
# amp_enabled=False,实际仍以 FP32 运行。
with torch.autocast(
device_type="cuda", dtype=amp_dtype, enabled=amp_enabled
):
logits = model(features) # 兼容的 CUDA 算子会以 BF16 计算。
predictions = logits.argmax(dim=1)
correct += (predictions == labels).sum().item()
total += labels.size(0)
return correct / total
def main():
train_dataset, valid_dataset = make_dataset()
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False)
# AMP 不要求将模型整体转为 BF16;模型参数默认保留 float32,
# autocast 只在执行兼容算子时临时选择 BF16 计算。
model = Classifier().to(device)
loss_fn = nn.CrossEntropyLoss() # 交叉熵会按 autocast 规则使用适合的数值精度。
optimizer = torch.optim.Adam(model.parameters(), lr=0.03) # Adam 更新仍保存为 FP32 的模型参数。
# 本例不创建 GradScaler。BF16 的指数范围接近 float32,通常不易因梯度过小而下溢;
# 因此直接调用 loss.backward() 即可。注意:“通常不需要”不代表绝不会出现 Inf/NaN,
# 学习率过高、数据异常或模型本身不稳定时仍可能产生非有限数值。
print(f"运行设备: {device}")
print(f"CUDA 是否支持 BF16: {bf16_supported}")
print(f"是否启用 CUDA BF16 AMP: {amp_enabled}")
epochs = 30
for epoch in range(epochs):
model.train()
total_loss = 0.0
total_samples = 0
for features, labels in train_loader:
features = features.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# autocast 仅包裹前向传播和损失计算,反向传播与 optimizer.step() 放在其外部。
#
# amp_enabled 是总开关:
# - 纯 CPU 环境中,device.type == "cpu",所以 amp_enabled=False。
# 此时 enabled=False 会禁用 autocast;这个 with 语句只是一个“不做任何
# dtype 转换”的上下文,features、模型参数、logits、loss 和梯度都保持
# 默认的 float32,训练过程就是普通 FP32 训练。
# - 支持 BF16 的 CUDA 环境中,amp_enabled=True。autocast 不会先把 features
# 或整个 model 永久转换为 BF16;它会在执行每个算子时自行决定计算 dtype。
# 通常卷积、矩阵乘法等兼容算子会使用 BF16,而归约等数值敏感算子会保持
# float32。因此,某个中间张量的 dtype 取决于产生它的算子。
with torch.autocast(
device_type="cuda", dtype=amp_dtype, enabled=amp_enabled
):
logits = model(features) # 在 CUDA BF16 AMP 时,模型内部的兼容算子会自动使用 BF16。
loss = loss_fn(logits, labels) # 损失函数内部的数值敏感部分会按 autocast 规则保持合适精度。
# BF16 通常不需要梯度缩放,因此直接根据当前 loss 计算梯度。
# 参数本身是 float32,所以 parameter.grad 通常也会以 float32 保存。
loss.backward()
# 使用刚计算出的梯度更新模型中保持为 float32 的参数。
# 与 FP16 的 GradScaler.step() 不同,这里不会自动检查 Inf/NaN 或跳过更新。
optimizer.step()
batch_size = labels.size(0)
total_loss += loss.item() * batch_size
total_samples += batch_size
train_loss = total_loss / total_samples
valid_accuracy = evaluate(model, valid_loader)
if (epoch + 1) % 5 == 0 or epoch == 0:
print(
f"Epoch [{epoch + 1:2d}/{epochs}] "
f"train_loss: {train_loss:.4f} "
f"valid_accuracy: {valid_accuracy:.2%}"
)
if __name__ == "__main__":
main()
"""
### 支持 BF16 的 NVIDIA 服务器上某次执行:
(py31213) root@controller01:~# python 52-amp混合精度训练-bf16.py
运行设备: cuda
CUDA 是否支持 BF16: True
是否启用 CUDA BF16 AMP: True
Epoch [ 1/30] train_loss: 0.4509 valid_accuracy: 85.00%
Epoch [ 5/30] train_loss: 0.2467 valid_accuracy: 86.25%
Epoch [10/30] train_loss: 0.2395 valid_accuracy: 83.75%
Epoch [15/30] train_loss: 0.2355 valid_accuracy: 87.50%
Epoch [20/30] train_loss: 0.2364 valid_accuracy: 87.50%
Epoch [25/30] train_loss: 0.2258 valid_accuracy: 87.50%
Epoch [30/30] train_loss: 0.2200 valid_accuracy: 85.00%
"""
6.6 数据增强
6.6.1 数据增强的背景与必要性
训练集数量有限时,模型可能记住训练图片的背景、拍摄角度或噪声,而不是学到类别真正具有的特征。这种“训练集表现很好、未见过的数据表现变差”的现象称为过拟合。收集更多真实且正确标注的数据通常最有效,但在医学图像、工业缺陷检测等场景中,采集和标注成本很高。
数据增强(data augmentation)是在训练时对已有样本施加合理变换,从一张训练图片生成不同但语义相同的训练输入。例如,一张猫的照片经轻微裁剪、左右翻转或亮度变化后,标签仍然是“猫”。这样可增加训练输入的多样性,使模型不那么依赖某个固定位置、亮度或背景。
数据增强并不是越强越好。变换必须符合任务的真实规律:自然图像分类中的轻微裁剪、翻转和颜色扰动通常合理;对“左右”本身有语义的任务、文字识别或医学影像,随意翻转、旋转或改变颜色可能改变标签含义。验证集与测试集用于模拟真实未见数据,通常只做确定性的尺寸调整与归一化,不应使用随机增强。
旧教程常使用 imgaug。它的维护已长期停滞,在 Python 3.12
/ NumPy 2 等新环境中可能难以安装或兼容。当前 PyTorch
图像任务中,推荐优先使用与 PyTorch 紧密集成的
torchvision.transforms.v2;当任务包含目标检测框、分割掩码、关键点等多种标注,或需要更丰富的增强算子时,Albumentations
也是常见选择。
6.6.2 torchvision.transforms.v2 的安装与使用
6.6.2.1 torchvision.transforms.v2 的介绍与安装
torchvision.transforms.v2 是 torchvision
提供的新一代图像变换 API。它可以处理 PIL 图片、PyTorch Tensor 以及
torchvision 的
TVTensor;对于图像分类,最常用的方式是把若干变换按顺序放进
v2.Compose。与旧版 torchvision.transforms
相比,v2 对 Tensor 和检测/分割任务的联合变换支持更好。
官方文档:https://docs.pytorch.org/vision/stable/transforms.html
。
使用前应保证 PyTorch 与 torchvision 的版本彼此匹配。以本教程的 PyTorch 2.6.1 为例,可安装与之对应的 torchvision 0.21.x;具体安装命令应根据操作系统、Python 版本及 CPU/CUDA 版本从 PyTorch 官网 生成。已安装后可验证:
1 | |
不要只为了使用 v2 而单独安装一个不匹配版本的 torchvision;版本不匹配可能导致导入错误或 C++ 算子不可用。
6.6.2.2 torchvision.transforms.v2 的使用
下面是图像分类训练集的变换管线。它将 PIL 的 RGB
图片随机裁剪到 224×224、随机水平翻转,然后转换成范围为
[0, 1] 的 float32 Tensor 并按 ImageNet
的均值/标准差归一化。最后输出形状为 (3, 224, 224) 的
Tensor:
1 | |
若使用 torchvision.datasets.ImageFolder,将训练变换传给
transform 参数即可;DataLoader 每次取样时会自动调用它:
1 | |
对于语义分割、目标检测等任务,不能只变换图像而不变换标注。v2
支持把图片与
tv_tensors.Mask、tv_tensors.BoundingBoxes
一起传给同一个几何变换,使随机裁剪、翻转的空间位置同步;具体可参考 torchvision v2
文档。
6.6.3 Albumentations 的安装与使用
6.6.3.1 Albumentations 的介绍与安装
Albumentations
是面向图像任务的数据增强库,提供大量图像、几何和天气/噪声等变换,并能把同一个随机几何变换同步应用到图像、分割掩码、检测框和关键点。它的输入通常是
NumPy 数组,图像形状为 (H, W, C);使用
ToTensorV2 后,输出才变为 PyTorch 的 (C, H, W)
Tensor。
官网:https://albumentations.ai,官方文档:https://albumentations.ai/docs/
。
在本教程的 ai_py312 虚拟环境中安装:
1 | |
Albumentations 会安装其需要的依赖。若运行图像读写或特定 OpenCV
变换时提示缺少 cv2,可按实际环境安装
opencv-python-headless(服务器、无 GUI 环境)或
opencv-python(桌面环境);两者通常只保留一个:
1 | |
6.6.3.2 Albumentations 的使用
6.6.3.2.1 Albumentations 的
Compose 简单操作
下面仍以图像分类为例。Albumentations 的 Compose
返回一个字典,图片位于 result["image"]
中。ToTensorV2 必须放在变换序列最后,因为前面的
Albumentations 变换处理的是 NumPy 图像:
1 | |
6.6.3.2.2 分割任务时的单张变换
分割任务应在 Compose
调用时同时传入掩码,保证随机几何变换同步:
1 | |
如下给出一个二分类分割任务的相对完整可执行示例:
1 | |
6.6.3.2.3 分割任务时的批量变换
分割任务需要返回逐像素对齐的
(image, mask),并把两者一起送进
transform(image=..., mask=...)。对于这类稍复杂的获取数据并进行数据增强的操作,不能直接使用ImageFolder,它是面向“整张图片一个类别标签”的分类任务。
ImageFolder 的单个样本返回
(image, class_index),其中 class_index
是目录名称对应的一个整数,不能表示整张逐像素掩码。也不能把 image 与 mask
分别交给独立的 transform /
target_transform:两个随机变换会分别抽取随机参数,可能使图片与掩码错位。
正确做法是自定义 Dataset。每次
__getitem__() 读取一对 image/mask,并在同一次
transform(image=..., mask=...) 调用中完成同步增强,返回
(image_tensor, mask_tensor);随后由 DataLoader
将多个已对齐的单样本堆叠成一个 batch。
本例沿用第 62 个文件的命名约定,目录可组织为:
1 | |
每张掩码必须与对应 RGB 图片宽高完全一致。本节为“目标 /
背景”的二分类分割示例:mask 使用单通道无损 PNG,背景像素为
0,目标像素为 255。不要用 JPG 保存
mask,因为有损压缩可能产生 1、254
等不属于类别的值。
1 | |
也就是说,Dataset.__getitem__()
负责“读取一对样本并同步增强”,DataLoader
负责“把多个已对齐样本组成 batch”。二分类分割模型通常输出
(batch_size, 1, 224, 224) 的原始 logits,可直接与上述
masks 配合 BCEWithLogitsLoss 使用。
多类别分割的 mask 不应使用 0/255,而应直接保存类别编号
0, 1, 2, ...;此时通常返回形状为
(batch_size, H, W) 的 long 掩码,并配合
CrossEntropyLoss。
使用增强时,先从任务语义允许的轻量变换开始,并在单独的验证集上比较增强前后的指标。若训练集准确率下降、验证集也无提升,或变换改变了标签含义,应减弱或删除该变换。
6.6.3.2.4 分类任务批处理
此处以每张图片只属于某个分类集合中唯一一类为例进行说明。此类任务中,对于每张图片有一个类别标签,可以直接使用ImageFolder来遍历按类别存放的图片,它会为每个类别分配一个类别编号,比如
0 表示 audi、1 表示
bmw。虽然类别目录名称可以自行替换,但训练集与验证集/测试集中的类别目录必须一致。
完整的示例代码如下: 1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247"""使用 ImageFolder 与 DataLoader 批量处理图像分类数据的示例。
本文件说明“整张图片一个类别标签”的分类任务如何进行批处理:
磁盘中的单张图片
-> ImageFolder 读取图片、根据所在类别目录生成 class_index
-> transform 处理单张图片
-> DataLoader 将多个 (image, class_index) 堆叠成一个 batch
-> 模型一次处理整个 batch,计算分类损失并更新参数
本示例是分类任务,不是分割任务:
- 分类的单样本标签是一个整数,例如 0 表示 audi、1 表示 bmw;
- 分割的单样本标签是一张与图片逐像素对应的 mask,不能使用 ImageFolder。
运行前请按下面结构准备图片。类别目录名称可以自行替换,但 train 与 eval 中的
类别目录必须一致:
data/
├── train/
│ ├── audi/
│ │ ├── audi_001.jpg
│ │ └── audi_002.jpg
│ └── bmw/
│ ├── bmw_001.jpg
│ └── bmw_002.jpg
└── eval/
├── audi/
│ └── audi_003.jpg
└── bmw/
└── bmw_003.jpg
ImageFolder 会按类别目录名的字母顺序自动分配类别编号。在上面的例子中通常是:
{"audi": 0, "bmw": 1}
实际映射以程序打印出的 dataset.class_to_idx 为准。
"""
from pathlib import Path # 用于检查训练集与验证集目录是否存在。
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import v2
# 固定随机种子,使模型初始化和 DataLoader 打乱顺序在大多数情况下可复现。
torch.manual_seed(42)
# 有 CUDA GPU 时在 GPU 训练,否则自动使用 CPU。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 将目录统一写成常量。若数据放在其他位置,只需修改这里。
TRAIN_DIR = Path("../stanford-cars-dataset-data/train")
EVAL_DIR = Path("../stanford-cars-dataset-data/eval")
# ImageFolder 的 transform 参数接收“单张图片”,并在每一次 __getitem__ 时调用。
# 训练集使用随机裁剪、随机翻转,使同一张训练图片在不同 epoch 中呈现略有不同的样子。
train_transform = v2.Compose([
v2.ToImage(), # PIL.Image -> 图像 Tensor,形状从 (H, W, C) 概念转换为 (C, H, W)。
v2.RandomResizedCrop(
size=(224, 224), # 模型最终接收的图片尺寸。
scale=(0.8, 1.0), # 随机裁取原图面积的 80%~100%,再缩放到 224×224。
antialias=True,
),
v2.RandomHorizontalFlip(p=0.5), # 每张训练图片有 50% 概率水平镜像翻转。
v2.ToDtype(torch.float32, scale=True), # uint8 [0, 255] -> float32 [0, 1]。
v2.Normalize(
mean=(0.485, 0.456, 0.406), # ImageNet RGB 通道均值。
std=(0.229, 0.224, 0.225), # ImageNet RGB 通道标准差。
),
])
# 验证/测试集不能使用随机增强,否则每次评估输入都不同,准确率会额外波动。
eval_transform = v2.Compose([
v2.ToImage(),
v2.Resize(size=(256, 256), antialias=True), # 固定缩放整张图。
v2.CenterCrop(size=(224, 224)), # 固定从中心裁取 224×224。
v2.ToDtype(torch.float32, scale=True),
v2.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
),
])
class SmallClassifier(nn.Module):
"""用于演示 batch 输入/输出形状的小型分类网络。
输入 images 的形状为 (B, 3, 224, 224),其中 B 是 batch_size;
输出 logits 的形状为 (B, num_classes),每一行是该图片属于各类别的原始分数。
"""
def __init__(self, num_classes):
super().__init__()
# 特征提取部分:逐步降低图片空间尺寸,并增加通道数。
self.features = nn.Sequential(
nn.Conv2d(3, 16, kernel_size=3, padding=1), # (B, 3, 224, 224) -> (B, 16, 224, 224)
nn.ReLU(),
nn.MaxPool2d(kernel_size=2), # -> (B, 16, 112, 112)
nn.Conv2d(16, 32, kernel_size=3, padding=1), # -> (B, 32, 112, 112)
nn.ReLU(),
nn.MaxPool2d(kernel_size=2), # -> (B, 32, 56, 56)
nn.AdaptiveAvgPool2d(output_size=1), # 每个通道压缩为一个数:-> (B, 32, 1, 1)
)
# 分类器:Flatten 将 (B, 32, 1, 1) 展平为 (B, 32),Linear 输出类别 logits。
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(32, num_classes),
)
def forward(self, images):
features = self.features(images)
return self.classifier(features)
def evaluate(model, data_loader):
"""在验证集上计算准确率;评估阶段不计算梯度、不更新参数。"""
model.eval() # 切换评估模式;Dropout/BatchNorm 等层会采用评估行为。
correct = 0
total = 0
with torch.no_grad():
# DataLoader 每次给出一个 batch:
# images 形状为 (当前 batch 实际大小, 3, 224, 224);
# labels 形状为 (当前 batch 实际大小,)。
for images, labels in data_loader:
images = images.to(device)
labels = labels.to(device)
logits = model(images) # 两类别时形状为 (B, 2)。
predictions = logits.argmax(dim=1) # 每行最大 logit 的下标,即预测类别编号。
correct += (predictions == labels).sum().item()
total += labels.size(0)
return correct / total
def main():
# 在构造 ImageFolder 前检查目录。未准备数据时给出清楚提示,而不是产生难理解的内部报错。
for directory in (TRAIN_DIR, EVAL_DIR):
if not directory.is_dir():
raise FileNotFoundError(
f"找不到数据目录:{directory}\\n"
"请按本文件顶部的 data/train/类别名、data/eval/类别名 结构放入图片,"
"或修改 TRAIN_DIR、EVAL_DIR。"
)
# 创建 ImageFolder 时,先扫描类别子目录名并建立类别编号,例如:
# {"audi": 0, "bmw": 1}。它还会在内部的 samples 列表中记录每个样本的
# (图片文件路径, 类别编号) 配对,例如 ("data/train/audi/00001.jpg", 0)。
# 注意:此时不会读取所有图片,更不会提前把所有图片转换为 (3, 224, 224)
# 的 Tensor 存入内存;这里只是保存“图片在哪里、它属于哪个类别”的索引。
# 真正在代码其他位置执行 train_dataset[index] 时,ImageFolder 才会:
# 1. 根据 samples[index] 从磁盘打开一张图片;
# 2. 调用这里传入的 train_transform;
# 3. 得到形状为 (3, 224, 224) 的图像 Tensor,并返回 (image_tensor, class_index)。
# 因此,训练变换中的随机裁剪/翻转会在每次取样时重新发生;同一张图片的
# Tensor 内容可随 epoch 改变,但它配对的 class_index 始终不变。
train_dataset = ImageFolder(TRAIN_DIR, transform=train_transform)
# 验证集的建立机制完全相同;区别仅在于传入 eval_transform,其中没有随机增强。
eval_dataset = ImageFolder(EVAL_DIR, transform=eval_transform)
# train / eval 的类别映射必须完全相同,否则标签编号的含义会错位。
if train_dataset.class_to_idx != eval_dataset.class_to_idx:
raise ValueError(
"训练集与验证集的类别目录不一致:\\n"
f"train: {train_dataset.class_to_idx}\\n"
f"eval: {eval_dataset.class_to_idx}"
)
print("类别编号映射:", train_dataset.class_to_idx)
print("训练图片数量:", len(train_dataset))
print("验证图片数量:", len(eval_dataset))
# DataLoader 负责按索引从 train_dataset 取多个单样本,并自动堆叠成 batch。
# shuffle=True 并不是把磁盘中的图片文件或 labels 分别打乱,而是在每个 epoch
# 为“样本索引”生成一个随机顺序。每次取到一个索引时,ImageFolder 仍返回该索引
# 原本绑定的 (image_tensor, class_index),所以图片与标签绝不会因为打乱而错配。
# 例如本批随机选中 8 个索引时,DataLoader 会得到 8 个 (图片, 标签) 配对,再合并为:
# images.shape == (8, 3, 224, 224),labels.shape == (8,)。
# shuffle=False:验证时保持固定索引顺序,方便让每次评估使用相同的图片顺序。
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
eval_loader = DataLoader(eval_dataset, batch_size=8, shuffle=False)
# num_classes 不要手写为 2,而是依据实际类别目录的数量自动确定。
model = SmallClassifier(num_classes=len(train_dataset.classes)).to(device)
loss_fn = nn.CrossEntropyLoss() # 接收 (B, 类别数) logits 与 (B,) 整数类别标签。
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
print(f"运行设备:{device}")
epochs = 5
for epoch in range(epochs):
model.train()
loss_sum = 0.0
sample_count = 0
# DataLoader 已经完成批处理。若 batch_size=8,通常:
# images.shape == (8, 3, 224, 224)
# labels.shape == (8,)
# 最后一个 batch 的实际图片数量可能不足 8。
for images, labels in train_loader:
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
logits = model(images) # 一次同时处理当前 batch 的全部图片。
loss = loss_fn(logits, labels) # 默认是当前 batch 内所有样本损失的平均值。
loss.backward()
optimizer.step()
# 乘以当前实际 batch 大小后累加,避免最后一个小 batch 影响 epoch 平均损失。
current_batch_size = labels.size(0)
loss_sum += loss.item() * current_batch_size
sample_count += current_batch_size
train_loss = loss_sum / sample_count
eval_accuracy = evaluate(model, eval_loader)
print(
f"Epoch [{epoch + 1}/{epochs}] "
f"train_loss: {train_loss:.4f} "
f"eval_accuracy: {eval_accuracy:.2%}"
)
if __name__ == "__main__":
main()
"""
### 某次执行后的输出如下(暂不重点关注损失值大小与验证准确率高低):
(ai_py312) jsy@jmac 06第六章 % python 64-数据增强albumentations的使用-分类任务批处理.py
类别编号映射: {'acura': 0, 'am_general': 1, 'aston_martin': 2, 'audi': 3, 'bentley': 4, 'bmw': 5, 'bugatti': 6, 'buick': 7, 'cadillac': 8, 'chevrolet': 9, 'chrysler': 10, 'daewoo': 11, 'dodge': 12, 'eagle': 13, 'ferrari': 14, 'fiat': 15, 'fisker': 16, 'ford': 17, 'geo': 18, 'gmc': 19, 'honda': 20, 'hummer': 21, 'hyundai': 22, 'infiniti': 23, 'isuzu': 24, 'jaguar': 25, 'jeep': 26, 'lamborghini': 27, 'land_rover': 28, 'lincoln': 29, 'maybach': 30, 'mazda': 31, 'mclaren': 32, 'mercedes_benz': 33, 'mini': 34, 'mitsubishi': 35, 'nissan': 36, 'plymouth': 37, 'porsche': 38, 'ram': 39, 'rolls_royce': 40, 'scion': 41, 'smart': 42, 'spyker': 43, 'suzuki': 44, 'tesla': 45, 'toyota': 46, 'volkswagen': 47, 'volvo': 48}
训练图片数量: 6516
验证图片数量: 1628
运行设备:cpu
Epoch [1/5] train_loss: 3.4773 eval_accuracy: 10.81%
Epoch [2/5] train_loss: 3.4162 eval_accuracy: 10.75%
Epoch [3/5] train_loss: 3.4015 eval_accuracy: 11.12%
Epoch [4/5] train_loss: 3.3902 eval_accuracy: 10.93%
Epoch [5/5] train_loss: 3.3807 eval_accuracy: 11.00%
"""
此示例尝试训练对包含汽车的图片进行识别与分类,将每张图片划分到某个具体的汽车品牌分类上。其中使用的汽车图片数据来自“Stanford
Cars Dataset”(其他类似或相关的数据集有:“Comprehensive Cars
Dataset”、“Vehicle Make and Model Recognition
Database”),首先从kaggle官网此页面中下载:stanford-cars-dataset
,并经过一定处理后将 cars_train_annos.mat
中描述的8144张图片按照4:1分成两个部分:6516 张训练图片和 1628
张验证图片,一共涉及 到49
个汽车品牌目录。(原数据集文件中cars_test_annos.mat 不包含
class
车型字段,所以无法直接知道其具体汽车品牌分类,所以未在此处使用。)
6.7 使用argparse进行调参
本节学习的主要内容包括:深度学习的超参数的概念、argparse
的介绍与经常性用法,最后给出一个示例。
在学习argparse之前,先来学习或复习下深度学习中超参数是什么东西。
6.7.1 深度学习的超参数是什么
在深度学习中,超参数(Hyperparameter) 指的是在模型开始训练之前,由开发者手动设定并指定的参数配置。与之相对的概念是参数(Parameter)。以下给出了深度学习中超参数与参数简单比较:
| 概念 | 定义 | 谁来决定? | 典型例子 |
|---|---|---|---|
| 超参数 (Hyperparameter) | 控制模型结构或训练过程的外部配置项 | 人为手动设定(凭经验或调优算法) | 学习率、Batch Size、网络层数、Epoch 次数 |
| 参数 (Parameter) | 模型在训练过程中通过数据自动学习和更新出来的权重值 | 模型自己学(通过反向传播和梯度下降) | 神经元的权重 \(w\)、偏置 \(b\) |
6.7.1.1 经典超参数分类
超参数大致可以分为以下两类:
6.7.1.1.1 控制“训练过程”的超参数(优化超参数)
这类超参数决定了模型怎么学、学多久、学多快:
- 学习率(Learning Rate, \(lr\)):最重要的超参数之一。决定了每次更新权重时的步长大小。步长太大容易无法收敛(跳过极小值),步长太小训练速度太慢。
- 批次大小(Batch Size):每次送入网络进行计算和梯度更新的样本数量(如 32, 64, 128)。
- 训练轮数(Epochs):整个数据集被模型完整遍历训练的总次数。
- 优化器(Optimizer)选择:如 SGD、Adam、AdamW 等。
- 损失函数权重 / 正则化系数(Weight Decay / L2 Regularization):用于控制正则化强度,防止模型过拟合。
- 随机种子(Random Seed):控制权重初始化、数据 Shuffle 和 Dropout 等所有随机行为的起点,用于确保训练过程和实验结果的可复现性。
6.7.1.1.2 控制“网络结构”的超参数(结构超参数)
这类超参数决定了模型长什么样、容量多大:
- 网络层数(Depth):隐藏层的数量(如 ResNet-18 vs ResNet-50)。
- 每层的神经元数量 / 通道数(Width/Channels):例如卷积层输出的通道数(16, 32, 64...)。
- 激活函数(Activation Function):如 ReLU、GELU、LeakyReLU 的选择。
- Dropout 比例:训练时随机丢弃神经元的概率(如 0.2、0.5),用于防止过拟合。
6.7.1.2 为什么叫“超”参数?
这里的“超(Hyper-)”表示元(Meta)或更高层级的意思。
因为它们超越了模型自身的学习机制,是在定义“模型的学习规则”和“模型的容量边界”。模型在训练时,只能在其给定的超参数框架内去求解最优的参数(权重 \(w\) 和偏置 \(b\))。
6.7.1.3 常见调优方式(调参 / Hyperparameter Tuning)
找到一组优秀的超参数通常是深度学习实验中最耗费时间的工作之一,常见的调优策略包括:
- 网格搜索(Grid Search):穷举法。在预先设定的参数组合列表里遍历测试所有可能(耗时极长)。
- 随机搜索(Random Search):在参数空间内随机抽样组合。通常比网格搜索更高效,能更快覆盖更大的搜索空间。
- 贝叶斯优化(Bayesian Optimization):根据已测试组合的结果,用概率模型预测下一组最有潜力的超参数组合(自动化调参库如 Optuna、Ray Tune 常用的核心策略)。
- 凭经验与手动调整(Intuition / Manual Tuning):根据训练/验证集的 Loss 曲线和评估指标,人工手动修正(例如 Loss 震荡不收敛则调小学习率,发生严重过拟合则加大 Dropout 或 Weight Decay)。
6.7.2 argparse是什么
如上所述,深度学习中超参数有很多,如何对它们进行有序管理与组织是个要关注的问题,因为超参数是深度学习实验中很重要但也也很耗时的工作,可能需要我们经过多次试验才能找到一组合适且高效的超参数组合。
大多数超参数需要在训练过程中根据训练结果进行尝试性改进,比如说学习率这个超参数(其他很多超参数也类似),可能需要在前一次或多次实验结果的基础上进行分析与修改(比如);另外个别超参数,如 随机种子 一般需要在训练过程中保持固定不变,以确保训练过程和实验结果的可复现性。
这一大批超参数如何传入、如何解析、有没有合理的提示、有没有进行类型等检查,最好能统一的管理起来。Python已经原生内置了argparse这个模块来完成上述参数的传入与管理工作(深度学习中的超参数本质上也是程序参数,刚好利用上它)。
注意,除了利用argparse来管理传入的超参数,有时为了保证训练的连贯性还需要进行超参数的保存。
argparse 模块让编写用户友好的命令行接口变得容易。
程序定义它需要哪些参数,argparse 将会知道如何从 sys.argv 解析它们。
argparse 模块还能自动生成帮助和用法消息文本。
该模块还会在用户向程序传入无效参数时发出错误提示。
Python官方针对argparse模块的官方文档:argparse.html
。
6.7.3
argparse的经常性用法
6.7.3.1 用法概述
argparse模块对命令行接口的支持是围绕 argparse.ArgumentParser
展开的,它是一个参数规格容器,而且包含应用于参数解析器上一组配置选项。
使用argparse模块进行参数的提示、接收、解析与使用,可以归纳如下几个步骤:
- 创建
argparse.ArgumentParser实例。它是argparse模块工作的基础与入口。 - 执行
ArgumentParser.add_argument(xxx)方法,向argparse.ArgumentParser实例中添加参数规格,这个步骤规定了各个参数的名字、类型、描述、是否必传等规格与限制。它支持位置参数与非位置参数两种参数传入模式。 - 执行
ArgumentParser.parse_args()方法运行解析器,解析参数,并将提取的数据放入一个argparse.Namespace对象 - 通过点运算符(如
args.lr)获取解析到的参数值,使用。
一般来说,上述几个步骤中用法稍微复杂点的地方是上述第二个步骤,即”向argparse.ArgumentParser实例中添加参数规格“这个步骤,因为其中涉及的常用参数稍显多点。第一个步骤”创建argparse.ArgumentParser实例“时虽然也支持多个不同参数,但一般常用的相对较少。
6.7.3.2
argparse.ArgumentParser常用参数
创建解析器的基本写法如下:
1 | |
ArgumentParser(...) 的参数用于配置“整个命令行程序”的行为;每一个具体超参数的名称、类型和默认值,则在稍后的 parser.add_argument(...) 中定义。下表列出深度学习训练脚本中最常见的解析器参数:
| 参数 | 作用 | 常见用法 |
|---|---|---|
| description | 程序的一句话说明;用户执行 --help 时会显示它。 | description="训练汽车品牌分类模型。" |
| prog | 帮助文本中显示的程序名。默认取执行时的脚本名;通常可不写。 | prog="train.py" |
| epilog | 帮助文本末尾的补充说明,例如数据目录结构、示例命令。 | epilog="示例:python train.py --epochs 20" |
| formatter_class | 控制帮助文本的排版方式。训练脚本最常用的是 argparse.ArgumentDefaultsHelpFormatter,它会在帮助信息中自动显示默认值。 | formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| add_help | 是否自动添加 -h/--help。默认是 True,一般保持默认即可。 | add_help=True |
例如,下面的解析器会让 --help 输出中自动附带每个参数的默认值:
1 | |
此时执行:
1 | |
不会开始训练,而是显示用法、description、后续通过 add_argument() 注册的参数规格。这是检查“参数名是否写对、默认值是什么”的首选方法。
6.7.3.3
ArgumentParser.add_argument常用参数
6.7.3.3.1 常见参数概述
以下是执行ArgumentParser.add_argument(xxx)添加参数规格时经常使用到的参数及其解释:
先看一个典型的训练超参数:
1 | |
它规定了用户可在命令行传入
--lr 0.0003;解析后得到的值会自动转为 float,可通过 args.lr
获取。若用户不传入 --lr 参数及其值,则 args.lr 为默认值
1e-3。
| 参数 | 作用与要点 | 深度学习中的例子 |
|---|---|---|
| name_or_flags | add_argument() 的第一个或多个位置参数,用来指定命令行名称。名称不以 - 开头的是位置参数,必须按位置传入;以 -- 开头的是可选参数,通常用来表示超参数。长选项中的连字符会自动转为下划线属性名。 | "data_dir" → args.data_dir;"--batch-size" → args.batch_size。 |
| type | 将命令行中原本的字符串转换为指定类型。若不能转换,解析器会报错并退出。 | type=int 用于 --epochs、--batch-size;type=float 用于 --lr、--weight-decay;type=Path 用于目录路径。 |
| default | 用户没有传入该可选参数时使用的默认值。默认值应与 type 对应。 | default=32、default=1e-3、default="auto"。 |
| help | --help 中显示的说明文字。应写清参数含义、单位或取值范围,而不仅是重复参数名。 | help="每次参数更新使用的样本数。" |
| choices | 限制允许的值。传入不在候选集合内的值会立即报错,有助于避免拼写错误。 | choices=["sgd", "adamw"];choices=["auto", "cpu", "cuda"]。 |
| required | 对可选参数(--...)指定是否必须传入。位置参数本身默认必传,不需要设置它。路径、实验名称等没有合理默认值时可使用。 | required=True。训练超参数通常有默认值,不宜滥用。 |
| action | 指定“见到该参数时做什么”。布尔开关不要写成 type=bool,而应使用 store_true 或 store_false。 | action="store_true":写 --amp 时 args.amp=True;不写时为 False。 |
| dest | 指定解析结果在 args 中的属性名。一般不需要写,因为会由长选项自动推导。 | "--batch-size" 默认生成 args.batch_size;必要时可写 dest="batch_size"。 |
| metavar | 只改变帮助文本中“参数值占位符”的显示名,不改变解析结果。 | metavar="DIR",帮助中显示 --data-dir DIR。 |
| nargs | 指定一个参数接收几个值。只有确实需要输入多个值时才使用;单值超参数通常不需要。 | nargs=2, type=int 可接收 --image-size 224 224。 |
位置参数与可选参数的对比如下:
1 | |
布尔参数是初学者容易出错的地方。不要写
type=bool:命令行中的 "False"
是一个非空字符串,转换时仍可能得到
True。更清晰的写法是:
1 | |
上例中,未传 --no-augment 时 args.augment
为 True;传入该选项后才变为 False。
6.7.3.3.2 重点讲述action与default参数
在 ArgumentParser.add_argument() 中同时传入
action 与 default 参数时,行为完全取决于
action 的具体类型。
最常见的几种情况如下:
1.
action="store_true" 或
action="store_false"
这是布尔开关。是最常见同时使用 default 的场景。
- 原理:
- 当用户在命令行显式传入该参数标志时,触发指定的动作(设为
True或False)。 - 当用户在命令行未传入该参数标志时,直接使用
default指定的值。
- 当用户在命令行显式传入该参数标志时,触发指定的动作(设为
1 | |
注意:如果不显式写
default,store_true默认的default本身就是False,store_false默认的default本身就是True。但显式指定default能让代码意图更清晰,或者在特殊需求下反向定义(例如将action="store_true"的default设为None,以区分“用户未传”与“用户显式关闭/开启”)。
2.
action="store"(默认动作,存入单一值)
- 原理:
store是默认动作。如果命令行指定了该参数,则保存用户给定的值;如果命令行未指定,则保存default的值。
1 | |
3.
action="append"(追加到列表)
- 原理:每次在命令行遇到该参数时,将值追加(append)到列表中。
- 与 default 的冲突点/易错点:
- 未指定参数:变量值为
default的原始值(建议设为列表或None)。 - 指定参数:用户传入的值会追加在
default列表的末尾,而不会覆盖default列表!
- 未指定参数:变量值为
1 | |
最佳实践:如果使用
action="append",通常建议设置default=None或在解析后再做判定,避免默认列表被污染。
4.
action="count"(统计出现次数)
- 原理:统计标志在命令行出现的次数(常用于日志详细程度
--verbose或-vvv)。 - 行为:
- 未传入时,变量值为
default。 - 传入 \(N\) 次时,变量值为
default + N(如果default为整型)或直接计数(如果default=None,首次出现从 1 开始计)。
- 未传入时,变量值为
1 | |
5.
action="store_const"(存储常数)
- 原理:当指定参数标志时,将
const参数赋予的值存入;未指定时,存入default的值。
1 | |
总结对照表
| action 类型 | 命令行未指定参数时 | 命令行指定参数时 | 注意事项 |
|---|---|---|---|
store (默认) |
结果为 default |
结果为 用户输入值 | 最标准的默认值覆盖策略 |
store_true |
结果为 default (默认 False) |
结果为 True |
常用作功能开关 |
store_false |
结果为 default (默认 True) |
结果为 False |
常用作禁用开关 |
append |
结果为 default |
在 default 列表后追加 |
若 default 是列表,不会被覆盖,只会追加 |
count |
结果为 default |
结果为 default + 出现次数 |
用于统计出现频次 |
store_const |
结果为 default |
结果为 const 值 |
必须配置 const 参数 |
6.7.3.4 argparse完整示例
以下是一个完整的使用上述常用参数的示例:
1 | |
将代码保存为 70-参数解析器经常性用法.py
后,可按下面方式运行:
1 | |
第二条命令解析后的部分结果等价于:
1 | |
命令行参数只负责“本次运行的配置入口”。为了让日后能够复现一次实验,实际训练时还应将 vars(args)(或 args 的内容)连同模型权重、训练日志一起保存到 JSON、YAML 或 checkpoint 中。
6.7.3.5 总结
可以看到,argparse模块为命令行的传入与解析提供了方便、实用的API接口,极大地方便深度学习中超参数的传入与解析。
一般地,为了复现或继承训练结果,深度学习会将超参数以json或dict或yaml等格式文件进行保存,模型结构与权重参数以pth等格式文件进行保存。