PyTorch06-PyTorch进阶训练技巧

六、PyTorch进阶训练技巧

6.1 自定义损失函数

尽管PyTorch官方针对主流模型为用户提供了很多内置损失函数,但随着应用场景的变化与自定义模型的出现,那些内置损失函数可能并不适合,于是产生了一些针对非通用模型的损失函数。有些损失函数是针对某特定场景使用,可以直接拿来使用或修改使用,有些场景或自定义模型甚至找不到现成或逼近的损失函数可用,此时就需要自定义损失函数。

本章节学习的内容包括:如何自定义损失函数。

6.1.1 以函数形式定义损失函数

损失函数其实就是用来计算模型输出数据与真实数据之间的差距的,所以计算损失的过程就是一个函数执行的过程。 所以理论上来说,可以通过定义一个函数的形式来自定义损失函数。

1
2
3
4
5
6
7
8
9
10
11
import torch
from torch import nn

def my_loss(output, target):
"""
最简单的损失函数示例:求取两个张量的均方误差,通常也称为 L2 损失(L2 Loss)
"""
# 求差->平方->求均值
loss = torch.mean((output - target)**2)
return loss

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
8
import 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\):代表预测区域与真实区域没有任何重合。
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
2
3
4
5
6
7
8
9
10
11
12
import torch

def dice_loss(input, target, smooth=1e-5):
# input: 模型预测概率 (N, C, H, W)
# target: 真实标签 (N, C, H, W)

# 将 Tensor 展平计算全局重叠
intersection = torch.sum(input * target)
cardinality = torch.sum(input) + torch.sum(target)

dsc = (2.0 * intersection + smooth) / (cardinality + smooth)
return 1.0 - dsc
6.1.2.2.4 为什么分割任务中非常喜欢用 Dice Loss?

普通的交叉熵损失(Cross Entropy Loss)是对所有像素一视同仁地计算误差。如果在医学图像中,病灶(前景目标)只占全图的 1%,背景占了 99%(严重的前背景极度不平衡),模型即使把所有像素都预测为背景,准确率也能达到 99%,但这种模型是无效的。

Dice Loss 关注的是前景与真实区域的交叉重叠比例,它天然消除了背景像素占绝大多数带来的影响,能够极大地改善类别不平衡(Class Imbalance)问题。

6.1.2.3 DiceBCELoss

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
import torch
import torch.nn as nn
import torch.nn.functional as F


class DiceBCELoss(nn.Module):
"""二分类或多标签分割的 BCE + Dice 组合损失。

inputs: 模型的原始输出 logits,形状为 (N, C, ...),不要预先经过 sigmoid。
targets: 与 inputs 形状相同、元素为 0 或 1 的标签。

BCE 逐像素衡量预测与标签的差异;Dice 关注预测区域和真实区域的重叠。
将两者相加常用于前景像素很少的二分类分割任务。
"""

def __init__(self, smooth=1e-6, bce_weight=1.0, dice_weight=1.0, reduction="mean"):
super().__init__()
if reduction not in {"mean", "sum", "none"}:
raise ValueError("reduction 必须是 'mean'、'sum' 或 'none'。")
self.smooth = smooth
self.bce_weight = bce_weight
self.dice_weight = dice_weight
self.reduction = reduction

def forward(self, inputs, targets):
# 模型输出与标签必须逐元素对应,例如都为 (N, 1, H, W)。
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) 两维张量。")

# 标签常是 int64 的 0/1 张量,计算前需转为与 logits 相同的浮点类型和设备。
targets = targets.to(device=inputs.device, dtype=inputs.dtype)

# 1. BCE 部分:直接处理 logits,数值上比 sigmoid 后再算 BCE 更稳定。
# reduction='none' 保留每个元素的损失,以便和“每样本、每通道”的 Dice 对齐。
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")

# 2. Dice 部分:sigmoid 将 logits 转换为 [0, 1] 的前景预测概率。
probabilities = torch.sigmoid(inputs)
# 保留 batch 维 N 与通道维 C;只在 H、W 等空间维求和。
spatial_dims = tuple(range(2, inputs.ndim))
if spatial_dims:
bce = bce.mean(dim=spatial_dims) # 形状:(N, C)
intersection = (probabilities * targets).sum(dim=spatial_dims)
cardinality = probabilities.sum(dim=spatial_dims) + targets.sum(dim=spatial_dims)
else:
# 兼容普通二分类张量 (N, C),此时每个元素本身就是一个计算单位。
intersection = probabilities * targets
cardinality = probabilities + targets

dice = (2 * intersection + self.smooth) / (cardinality + self.smooth)
dice_loss = 1 - dice # 形状:(N, C)

# 可通过 bce_weight、dice_weight 调整两部分在总损失中的相对影响。
loss = self.bce_weight * bce + self.dice_weight * dice_loss

if self.reduction == "mean":
return loss.mean() # 常用:返回一个标量,供 loss.backward() 使用。
if self.reduction == "sum":
return loss.sum()
return loss # 形状:(N, C),供调用方自行加权或分析。


# ==================== 使用方法:二分类分割 ====================
torch.manual_seed(42)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# N=2:两张图像;C=1:每个像素只有前景/背景两种可能;H、W:图像尺寸。
# 用 randn() 模拟模型的 logits,它可以为任意实数;requires_grad=True 用于验证反向传播。
logits = torch.randn(2, 1, 64, 64, device=device, requires_grad=True)
# 真实标签必须与 logits 形状一致,且元素为 0 或 1。
targets = torch.randint(0, 2, (2, 1, 64, 64), device=device)

criterion = DiceBCELoss(bce_weight=1.0, dice_weight=1.0)
loss = criterion(logits, targets)
loss.backward()

print("运行设备:", device)
print(f"BCE + Dice 损失:{loss.item():.6f}")
print("logits 是否获得梯度:", logits.grad is not None)

# 注意:多类别互斥分割(一个像素只能属于 C 个类别中的一个)通常使用 CrossEntropyLoss,
# 并需要 softmax 版本的 Dice;不能直接使用此处基于 sigmoid 的二分类/多标签实现。

6.1.2.4 IoULoss

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
import torch
import torch.nn as nn


class IoULoss(nn.Module):
"""二分类或多标签分割的 IoU(Jaccard)损失。

inputs 是形状为 (N, C, ...) 的 logits;targets 与其形状相同,且元素为 0 或 1。
IoU = 交集 / 并集,损失定义为 1 - IoU。每个样本、每个通道独立计算后再规约。
"""

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 转成概率。targets 转为同设备、同浮点类型,便于后续相乘和求和。
probabilities = torch.sigmoid(inputs)
targets = targets.to(device=inputs.device, dtype=probabilities.dtype)

# 对每个样本、每个通道,在空间维(如 H、W)独立计算 IoU。
spatial_dims = tuple(range(2, inputs.ndim))
if spatial_dims:
intersection = (probabilities * targets).sum(dim=spatial_dims)
# 并集 = 预测区域 + 标签区域 - 重叠区域。
union = (probabilities + targets - probabilities * targets).sum(dim=spatial_dims)
else:
# 兼容普通二分类张量 (N, C)。
intersection = probabilities * targets
union = probabilities + targets - probabilities * targets

# smooth 避免预测与标签均为空时出现 0/0;其值应远小于正常的区域面积。
iou = (intersection + self.smooth) / (union + self.smooth)
loss = 1 - iou

if self.reduction == "mean":
return loss.mean()
if self.reduction == "sum":
return loss.sum()
return loss # 形状:(N, C)


# ==================== 使用方法:二分类分割 ====================
torch.manual_seed(42)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 模拟模型输出 logits 与二值分割标签。输入不是概率,因此可以取任意实数。
logits = torch.randn(2, 1, 64, 64, device=device, requires_grad=True)
targets = torch.randint(0, 2, (2, 1, 64, 64), device=device)

criterion = IoULoss()
loss = criterion(logits, targets)
loss.backward()

print("运行设备:", device)
print(f"IoU 损失:{loss.item():.6f}")
print("logits 是否获得梯度:", logits.grad is not None)

# 注意:IoU 损失衡量区域重叠,对小目标分割很有帮助;实践中也常与 BCE 损失组合使用,
# 以同时提供逐像素监督和整体区域重叠监督。

6.1.2.5 FocalLoss

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
import torch
import torch.nn as nn
import torch.nn.functional as F


class FocalLoss(nn.Module):
"""用于二分类或多标签任务的 Binary Focal Loss。

inputs: 模型输出的 logits,形状为 (N, C, ...),不要预先经过 sigmoid。
targets: 与 inputs 形状相同、元素为 0 或 1 的标签。

gamma 会降低“容易分类样本”的损失权重;alpha 用于平衡正、负样本。
它适合前景很少的分割任务或正负样本不均衡的多标签任务。
"""

def __init__(self, alpha=0.25, gamma=2.0, reduction="mean"):
super().__init__()
if alpha is not None and not 0 <= alpha <= 1:
raise ValueError("alpha 必须位于 [0, 1],或设为 None 以不使用类别平衡。")
if gamma < 0:
raise ValueError("gamma 必须大于或等于 0。")
if reduction not in {"mean", "sum", "none"}:
raise ValueError("reduction 必须是 'mean'、'sum' 或 'none'。")
self.alpha = alpha
self.gamma = gamma
self.reduction = reduction

def forward(self, inputs, targets):
if inputs.shape != targets.shape:
raise ValueError(
"inputs 与 targets 的形状必须相同,"
f"当前分别为 {tuple(inputs.shape)} 和 {tuple(targets.shape)}。"
)

targets = targets.to(device=inputs.device, dtype=inputs.dtype)

# BCEWithLogits 的实现直接处理 logits,数值上比 sigmoid 后再算 BCE 更稳定。
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
probabilities = torch.sigmoid(inputs)

# p_t 表示“模型给真实类别分配的概率”:
# 标签为 1 时取 p;标签为 0 时取 (1 - p)。p_t 越接近 1,样本越容易。
p_t = probabilities * targets + (1 - probabilities) * (1 - targets)
focal_weight = (1 - p_t).pow(self.gamma)

# alpha 为正类设置 alpha,为负类设置 (1 - alpha)。
# alpha=None 时不进行正负类别平衡,仅使用 gamma 的困难样本聚焦作用。
if self.alpha is not None:
alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
focal_weight = alpha_t * focal_weight

loss = focal_weight * bce

if self.reduction == "mean":
return loss.mean()
if self.reduction == "sum":
return loss.sum()
return loss # 与 inputs 相同形状,供调用方自行按样本或像素加权。


# ==================== 使用方法:类别不平衡的二分类分割 ====================
torch.manual_seed(42)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 用较少的正类像素模拟“前景稀少”的分割场景;logits 是模型原始输出。
logits = torch.randn(2, 1, 64, 64, device=device, requires_grad=True)
targets = (torch.rand(2, 1, 64, 64, device=device) < 0.1).long()

# gamma=2 是常用起点;alpha=0.25 表示正类基础权重为 0.25、负类为 0.75。
# alpha 的具体值应结合任务中正负样本比例和验证集结果调节,而非机械固定。
criterion = FocalLoss(alpha=0.25, gamma=2.0)
loss = criterion(logits, targets)
loss.backward()

print("运行设备:", device)
print(f"正类像素比例:{targets.float().mean().item():.4f}")
print(f"Focal Loss:{loss.item():.6f}")
print("logits 是否获得梯度:", logits.grad is not None)

# 注意:多类别互斥分类/分割(每个样本或像素只能属于一个类别)需要 softmax 版本的 Focal Loss,
# 通常以 (N, C, ...) logits 和类别索引标签作为输入,不能直接套用本二分类版本。

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
18
CLASSES
builtins.object
LRScheduler
ChainedScheduler
ConstantLR
CosineAnnealingLR
CosineAnnealingWarmRestarts
CyclicLR
ExponentialLR
LambdaLR
LinearLR
MultiStepLR
MultiplicativeLR
OneCycleLR
PolynomialLR
ReduceLROnPlateau
SequentialLR
StepLR

也可以直接在PyTorch官方文档中查看它们的详细介绍:PyTorch官网optim文档 。每个学习率调整策略类的详细文档如下:

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
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
"""使用自定义学习参数调整函数 adjust_learning_rate 动态调整学习率。

本例使用一份可控的二维线性回归数据,不需要下载数据集,直接运行即可。
重点是理解官方 scheduler 的基本使用顺序:
先由 optimizer 更新模型参数,再在一个 epoch 结束时调用 adjust_learning_rate 动态调整学习率。
"""

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()

default_lr = 0.1
# optimizer 负责根据反向传播得到的梯度更新模型参数。
optimizer = torch.optim.Adam(model.parameters(), lr = default_lr)

def adjust_learning_rate(optimizer, epoch):
"""
假设我们现在正在做实验,需要学习率每30轮下降为原来的1/10,
假设已有的官方API中没有符合我们需求的,那就需要自定义函数来实现学习率的改变。
"""
lr = default_lr * (0.1 ** ((epoch+1) // 30))
optimizer.param_groups[0]['lr'] = lr
# scheduler = torch.optim.lr_scheduler.StepLR(
# optimizer, step_size=20, gamma=0.5
# )

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 的参数更新)结束后,再更新学习率。
# 调用后,下一轮 epoch 会使用更新后的学习率。
adjust_learning_rate(optimizer, epoch)

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()

6.2.2.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
176
177
178
179
180
181
"""使用自定义学习率调度类动态调整学习率。

本例使用一份可控的二维线性回归数据,不需要下载数据集,直接运行即可。
重点是理解:除了直接编写 adjust_learning_rate 函数外,也可以将策略封装成
一个拥有 step() 方法的类。训练循环便可以像使用官方 scheduler 一样调用它。
"""

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


class CustomStepLR:
"""自定义的按固定轮数衰减学习率的调度类。

它实现的策略与第 21 个示例相同:每完成 step_size 个 epoch,
将初始学习率乘以 gamma 一次。

例如,初始学习率为 0.1、step_size=30、gamma=0.1 时:
第 1~30 轮使用 0.1;第 31~60 轮使用 0.01;
第 61~90 轮使用 0.001。
"""

def __init__(self, optimizer, step_size, gamma):
# 保存 optimizer,后续在 step() 中通过它访问并修改 param_groups。
self.optimizer = optimizer

# step_size 必须为正数;否则“每多少轮衰减一次”没有意义。
if step_size <= 0:
raise ValueError("step_size 必须大于 0")
self.step_size = step_size
self.gamma = gamma

# 记录已经完成的 epoch 数。开始训练前为 0。
self.completed_epochs = 0

# 一个 optimizer 可以包含多个参数组,例如模型不同部分使用不同初始学习率。
# 保存每一组的初始学习率,而不是只修改 param_groups[0]。
self.base_lrs = [param_group["lr"] for param_group in optimizer.param_groups]

def step(self):
"""在一个 epoch 结束后调用,计算并写入下一轮要使用的学习率。"""
self.completed_epochs += 1

# // 是整除:完成第 30 轮时,decay_times 为 1;完成第 60 轮时为 2。
decay_times = self.completed_epochs // self.step_size

# 使用初始学习率重新计算,避免反复乘法带来的累计误差。
for base_lr, param_group in zip(self.base_lrs, self.optimizer.param_groups):
param_group["lr"] = base_lr * (self.gamma ** decay_times)

def get_last_lr(self):
"""返回当前已写入 optimizer 的各参数组学习率,作用类似官方 scheduler。"""
return [param_group["lr"] for param_group in self.optimizer.param_groups]


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)

# 创建自定义调度器。它不是 torch.optim.lr_scheduler 中的官方类,
# 但同样将 optimizer 传入,并通过 scheduler.step() 更新学习率。
scheduler = CustomStepLR(optimizer, step_size=30, gamma=0.1)

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 的参数更新)结束后,再更新学习率。
# 调用后,下一轮 epoch 会使用更新后的学习率。
scheduler.step()

# get_last_lr() 返回列表;本例只有一个参数组,取第 0 个元素即可。
current_lr = scheduler.get_last_lr()[0]

# 每 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()

6.3 模型微调torchvision

本节学习的内容包括:模型微调的流程、PyTorch中内置的主要开源模型、如何训练模型的指定层。

6.3.1 微调相关背景介绍

开源社区已经为大家提供了很多针对各种场景下使用的神经网络,其中很多神经网络训练时采用了大规模、高质量的数据集,但它们终归只是考虑到了当时场景。大家在实际学习工作时,会遇到各种不同实际业务场景,经当时的数据集训练得到的神经网络一般情况下并不能直接应用。

对于一般小型团队或公司而言,项目中能够收集到的训练、测试数据集规模相对较小,相比而言开源社区很多结构复杂的大规模神经网络如ViT使用的数据集则大得多,比如ImageNet系统数据集至少是百万个样本数据。如果直接采用这些”开源神经网络的结构+自己的目标数据集“,很可能会造成过拟合,同时因为此时的目标数据集小太多可能训练得到的最终模型精度(如准确率、泛化能力等)达不到实用程度。 (当然如果某开源模型当时就是由小规模数据集或跟自己当前的目标数据集规模度差不多甚至更小训练得到的,理论上是可能直接拿此开源模型的结构从头开始训练,更新得到目标数据集上的全新权重参数,从而得到一个新的目标模型。)

为此,对于这类大规模神经网络,要想使用它们的结构(甚至进行少量改动),并利用自己当前已有的目标数据集训练、更新权重参数,有如下两种常见的方法:

  • 将目标数据集规模、质量扩大或提升至开源的大规模神经网络当时使用的数据集一致。但收集与标数据集是一个成本非常高的过程,其成本大小与数据集规模、质量呈正相关关系,所以这一过程意味更高的成本,这对于一般的小团队或公司是不可取的。
  • 使用应用迁移学习(transfer learning),将开源模型在源数据集上学习到的能力迁移到目标数据集上。它的机理是:开源模型使用的源数据集跟目标数据集可能很大程度并不一致或不是描述同类事物,但开源模型在源数据集能够学习到抽象通用的图像特征(比如图像的纹理、边缘、形状与颜色等),它们也适用于目标数据集。

模型方面,应用迁移学习的最常用、最典型应用场景就是模型微调(finetuning)。模型微调的大致意思是:复制某个开源模型(包括模型结构与权重参数),在目标数据集上重新进行训练、不断更新其模型权重参数,使其能够在目标数据集表现达到预期。就是说我们站在巨人的肩膀上,利用优秀开源模型的骨架,赋予其目标数据集的灵魂,使模型与数据完美融合获得预期的涌现能力。

掌握模型微调,可以快速验证问题、提升模型与数据融合能力。

6.3.2 微调在 PyTorch 中的两类典型操作

在实际编写 PyTorch 代码时,根据目标数据集的大小,模型微调通常分为两种具体策略:

  1. 特征提取(Feature Extraction / 冻结骨干):
    • 操作:把开源 Backbone 中除了最后输出层外的所有层参数全冻结(param.requires_grad = False),只训练我们自己新加的最后一层分类头(Linear)。
    • 适用:目标数据集非常非常小(比如只有几百张图)时。

具体描述此策略下模型微调的流程如下:

  • 1)获取源模型:在源数据集(一般是大规模开源数据集)上训练得到源模型。这个工作一般开源社区已经帮忙完成了。
  • 2)构建目标模型:复制源模型中除了最后的输出层外的模型结构与权重参数,并手动构建、添加一个输出与目标数据集一致的新输出层,新输出层的权重参数是随机初始化的。
  • 3)更新新输出层:将目标模型中新输出层前面所有的层的权重参数冻结,在目标数据集上进行训练、更新新输出层中的权重参数。
  • 4)测试与保存目标模型:当目标模型测试达到预期后,将其以文件的形式持久化保存下来以备后用。如果测试不通过,则重复2)、3)、4)这几个步骤。 迁移学习_源模型与目标模型
  1. 全量/部分微调(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
2
3
4
5
6
dic = torch.utils.model_zoo.load_url(url="https://download.pytorch.org/models/resnet18-f37072fd.pth",
model_dir="/Users/jsy/.cache/torch/hub/test/test2")
#print(dic)
resnet18 = models.resnet18(weights=None)
dict_state = torch.load("/Users/jsy/.cache/torch/hub/test/test2/resnet18-f37072fd.pth")
resnet18.load_state_dict(dict_state)

6.3.5 训练特定层

默认情况下,模型中具体某个模型块或层的权重或偏置的requires_grad属性值为True,表示其在损失反向传播的过程中会记录梯度、被优化器自动更新,这在模型从头开始训练或全量微调过程中是正常无需修改的。 但正如上面所述,如果我们基于某个大规模开源模型构造了目标模型进行微调,其最后输出层是重新构建的(新构建的输出层默认权重与偏置的requires_grad属性值默认为True)、前面的层结构和权重参数是继承自原开源模型。此时就需要将目标模型最后输出层前面的层的权重参数进行冻结,以保证在微调过程中只更新输出层的权重参数,此时就需要将前面的层的权重或偏置的requires_grad属性值强制更改为False,表示在微调中不更新它们。

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
import torchvision
import torchvision.models as models
import torch.nn as nn

def set_parameter_requires_grad(model, feature_extracting):
if feature_extracting:
for param in model.parameters():
param.requires_grad = False

# 实例化 resnet18 模型
model = models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)

# 构建目标模型
## 冻结model所有层的权重或偏置 参数
feature_extract = True
set_parameter_requires_grad(model, feature_extract)
## 修改模型:创建新的输出层直接覆盖原输出层
num_ftrs = model.fc.in_features
model.fc = nn.Linear(in_features=num_ftrs, out_features=4, bias=True)

# print(f"model.fc: {model.fc}")
# print("model.fc.weight.requires_grad:", model.fc.weight.requires_grad)
# print("model.fc.bias.requires_grad: ", model.fc.bias.requires_grad)

print("\n查看模型所有参数的requires_grad属性值:")
for name, param in model.named_parameters():
print(name, param.requires_grad)

此示例代码基于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 timm
(2)源码可编辑安装
1
2
git 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
2
import timm
import torch

6.4.2.1 查询可用模块与模型

查询可用模块与模型,分别使用list_modules()和list_models() 方法。

timm.list_modules()返回已注册模型所在的模块名列表,返回值是按字母排序的list[str]。 它不是Python模块对象或文件路径,通常用于为list_models()的module参数选择范围。

1
2
3
4
5
6
7
# 查看部分可供筛选的模块名,例如'resnet'、'vision_transformer'
modules = timm.list_modules()
print(modules[:10])

# 只查询resnet模块中已注册的模型名称
resnet_names = timm.list_models(module='resnet')
print(resnet_names[:10])

timm.list_models()只读取模型注册表,不会创建模型,也不会下载权重。 其返回值是按自然顺序排序的模型名称列表。

1
2
3
4
5
6
7
8
9
10
# 列出全部已注册、且未被标记为废弃的模型名称
all_names = timm.list_models()
print(len(all_names), all_names[:10])

# 按名称搜索。筛选规则来自fnmatch:*表示任意长度字符,?表示一个字符
resnet50_names = timm.list_models(filter='resnet50*')
print(resnet50_names)

# 多个filter取并集:匹配任意一个规则的模型都会保留
names = timm.list_models(filter=['resnet50*', 'mobilenetv3*'])

可以同时按模块和名称筛选;两项条件需要同时满足。多个module也是并集。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# 只在vision_transformer模块中查找名称以vit_开头的模型
vit_names = timm.list_models(
filter='vit_*',
module='vision_transformer',
)

# 在两个模块中查找名称中含tiny的模型
tiny_names = timm.list_models(
filter='*tiny*',
module=['vision_transformer', 'swin_transformer'],
)

# 从结果中排除不需要的名称;多个排除规则同样可以同时使用
names_without_clip = timm.list_models(
filter='resnet*',
exclude_filters=['*clip*', '*gn*'],
)

pretrained=True表示只保留具有有效预训练权重配置的模型。此时include_tags默认自动设为True, 所以结果会带有具体权重标签,例如resnet50.a1_in1k:前半部分是架构名,点号后的部分是权重配置标签。 列出模型不会下载这些权重。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 查询resnet50系列的每一套具体预训练权重
pretrained_names = timm.list_models(
filter='resnet50*',
pretrained=True,
)
print(pretrained_names[:5])
# 可能包含:['resnet50.a1_in1k', 'resnet50.a1h_in1k', ...]

# 仅确认哪些架构至少有一套预训练权重,不显示具体权重标签
pretrained_arch_names = timm.list_models(
filter='resnet50*',
pretrained=True,
include_tags=False,
)

timm.list_pretrained()是查询具体预训练权重名称的快捷方法,等价于 list_models(pretrained=True, include_tags=True, ...)。当需要排除权重名称时,建议把 exclude_filters写成列表,即使只有一个规则。

1
2
3
4
5
6
7
8
# 只列出resnet18系列的具体预训练权重配置
resnet18_weights = timm.list_pretrained(filter='resnet18*')

# 排除某个权重标签;使用列表可正确处理带标签的结果
resnet18_weights = timm.list_pretrained(
filter='resnet18*',
exclude_filters=['*.a1_in1k'],
)

若批量测试模型而不希望同一架构的兼容别名重复出现,可加入name_matches_cfg=True。 日常选择模型时通常不需要该参数。

1
2
3
4
5
# 只保留名称与默认预训练配置名称匹配的模型,减少别名造成的重复
canonical_names = timm.list_models(
filter='resnet*',
name_matches_cfg=True,
)

6.4.2.2 确认模型与预训练权重

确认模型与预训练权重分别使用方法:is_model()和is_model_pretrained()。

在把字符串形式的模型名称交给训练脚本前,可以先进行检查。is_model()检查架构是否已注册; 即使传入带权重标签的名称,它也会识别其中的架构部分。is_model_pretrained()则检查该名称 是否有有效的预训练权重配置。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
model_name = 'resnet18'

# 架构是否已在当前timm中注册
if not timm.is_model(model_name):
raise ValueError(f'当前timm版本中没有模型:{model_name}')

# 该架构是否至少具有一套可用的预训练权重配置
if timm.is_model_pretrained(model_name):
print(f'{model_name}支持预训练权重')
else:
print(f'{model_name}没有注册预训练权重')

# 带标签时,检查的是指定的那套权重配置
weight_name = 'resnet18.a1_in1k'
print(timm.is_model(weight_name)) # True:resnet18架构存在
print(timm.is_model_pretrained(weight_name)) # True:该权重配置存在

这两个检查只查询本地注册表;它们不能保证权重文件已经缓存到本地,也不能保证网络一定能够 下载权重。

6.4.2.3 查看预训练配置

查看预训练配置主要使用方法:get_pretrained_cfg()和get_pretrained_cfg_value()。

一套预训练权重除了参数文件外,还对应输入尺寸、归一化均值和标准差、裁剪比例、类别数等配置。 get_pretrained_cfg()返回一个PretrainedCfg对象;传入带标签的名称可以精确查询某一套权重。 它只返回配置副本,不下载权重。

1
2
3
4
5
6
7
8
9
10
# 先从当前安装版本中取得一个确实存在的具体权重名称
weight_name = timm.list_pretrained(filter='resnet18*')[0]

# cfg描述该权重训练和推理时应使用的关键设置
cfg = timm.get_pretrained_cfg(weight_name)
print('权重名称:', weight_name)
print('输入形状(C, H, W):', cfg.input_size)
print('归一化均值:', cfg.mean)
print('归一化标准差:', cfg.std)
print('训练类别数:', cfg.num_classes)

如果只需要一个字段,可使用get_pretrained_cfg_value()。字段不存在时它返回None, 适合在程序中提供兼容处理。

1
2
3
4
5
6
# 读取该权重建议使用的插值方式,例如'bicubic'
interpolation = timm.get_pretrained_cfg_value(weight_name, 'interpolation')

# 查询不存在的字段不会抛出属性错误,而是得到None
unknown_value = timm.get_pretrained_cfg_value(weight_name, 'not_a_real_key')
print(interpolation, unknown_value)

应使用list_pretrained()或list_models(pretrained=True)返回的名称来查询配置。 对于不存在的架构或错误的权重标签,配置查询会报错;这能及早发现名称拼写问题。

6.4.2.4 创建与修改模型

创建模型使用的方法是create_model() 。

timm.create_model()根据模型名称查找构造函数,并返回一个torch.nn.Module。 pretrained=False(默认值)只创建随机初始化的模型;pretrained=True才会加载该模型或指定标签 对应的预训练权重。第一次加载预训练权重时,可能需要从网络下载并缓存文件。

1
2
3
4
5
6
7
8
# 创建一个随机初始化的ImageNet分类模型
model = timm.create_model('resnet18', pretrained=False)

# 输入张量的形状为(N, C, H, W),这里N=2、C=3、H=W=224
import torch
images = torch.randn(2, 3, 224, 224)
logits = model(images)
print(logits.shape) # torch.Size([2, 1000]),默认是1000个ImageNet类别

迁移学习时,最常见的是指定预训练模型并修改分类类别数。下面的num_classes=4会创建一个 输出4类的新分类头;预训练的主干网络参数会被加载,而新分类头需要在自己的数据集上训练。

1
2
3
4
5
6
7
8
9
10
11
12
# 指定具体权重标签,避免“同一架构有多套权重”时的歧义
weight_name = timm.list_pretrained(filter='resnet18*')[0]

# pretrained=True会按weight_name加载权重;num_classes=4将分类头改为4类输出
model = timm.create_model(
weight_name, # 'resnet18.a1_in1k'
pretrained=True,
num_classes=4,
)

logits = model(images)
print(logits.shape) # torch.Size([2, 4])

若只希望按架构默认的预训练配置加载,也可以使用不带标签的名称:

1
2
3
4
5
6
# 由timm为resnet18选择默认预训练配置
model = timm.create_model(
'resnet18',
pretrained=True,
num_classes=4,
)

甚至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
10
model = 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
2
3
4
5
6
7
8
9
10
# 取得注册的resnet18构造函数;此时尚未创建模型
builder = timm.model_entrypoint('resnet18')
print(builder.__name__) # 通常输出:resnet18

# 调用构造函数后才得到nn.Module实例
model = builder(
pretrained=False, # 不加载预训练权重
num_classes=4, # 创建4分类模型
)
print(type(model))

综合来看,日常任务通常使用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
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
import timm
import torch

# 利用 timm 创建
# print(f"timm.list_pretrained(filter='resnet18*'): {timm.list_pretrained(filter='resnet18*')}")
weight_name = timm.list_pretrained(filter='resnet18*')[0]
resnet18 = timm.create_model(
weight_name, # 'resnet18.a1_in1k'
pretrained=True,
num_classes=4, # 将原模型中最后输出1000分类的输出层,修改为输出4分类的输出层
)

print(type(resnet18)) # <class 'timm.models.resnet.ResNet'>
print(isinstance(resnet18, torch.nn.Module)) # True

# 假设这里对最后的输出4个分类的输出层进行微调

# 将模型的权重以文件的形式保存到当前目录下
save_path = "./resnet18.pth"
torch.save(resnet18.state_dict(), save_path)

# 再次创建随机初始化权重的 resnet18.a1_in1k 模型实例
weight_name2 = weight_name
model = timm.create_model(
weight_name2,
pretrained=False,
num_classes=4,
)
model.load_state_dict(torch.load(save_path))

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 为什么需要混合精度?

  1. FP32(单精度浮点数)的局限:
    • 显存大、算力未跑满:FP32 占用 4 字节。当模型参数量达到数亿或数十亿时,FP32 会迅速榨干 GPU 显存。
  2. 纯 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 混合精度训练的工程优势

  1. 显存消耗显著降低(减少约 40% ~ 50%):
    • 占显存大头的数据(激活值、中间张量)全部变为 16 位,允许提升 1.5~2 倍的 batch_size。
  2. 训练计算速度大幅提升(通常 1.5x ~ 3x 加速):
    • 极大地缓解了 GPU 显存带宽瓶颈(Memory Bandwidth Bound),同时全面释放了 Tensor Core 的矩阵计算吞吐量。
  3. 模型最终精度几乎无损:
    • 由于关键的权重更新与损失计算依然运行在 FP32 精度下,模型的收敛曲线与最终准确率/Loss 与全 FP32 训练高度吻合。

6.5.3.5 PyTorch 代码实现范式

在 PyTorch 2.6 中,推荐使用 torch.autocast 与 torch.amp.GradScaler。下面只保留训练循环中与 CUDA FP16 AMP 直接相关的关键代码;无 CUDA 时通过 enabled=False 自动退化为普通 FP32 训练:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
amp_enabled = device.type == "cuda"
amp_dtype = torch.float16 # 改为 torch.bfloat16 时,通常不需要 GradScaler
model = model.to(device)
scaler_enabled = amp_enabled and amp_dtype == torch.float16
scaler = torch.amp.GradScaler("cuda", enabled=scaler_enabled)

for inputs, targets in train_loader:
inputs = inputs.to(device)
targets = targets.to(device)
optimizer.zero_grad()

# 1. 前向传播与损失计算:适合的 CUDA 算子自动使用 amp_dtype
with torch.autocast("cuda", dtype=amp_dtype, enabled=amp_enabled):
outputs = model(inputs)
loss = criterion(outputs, targets)

# 2. 反向传播与优化器更新
scaler.scale(loss).backward() # 放大 loss 后计算梯度
scaler.step(optimizer) # 还原梯度,检查溢出后再更新参数
scaler.update() # 动态调整缩放因子

以下是完整的可执行示例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
2
3
4
5
6
7
8
jsy@jmac ~ % conda run -n ai_py312 python -c "import torch, torchvision; print(torch.__version__); print(torchvision.__version__)"
# 输出如下:
2.6.0
0.21.0

jsy@jmac ~ % conda run -n ai_py312 python -c "from torchvision.transforms import v2; print(v2.Compose)"
# 输出如下:
<class 'torchvision.transforms.v2._container.Compose'>

不要只为了使用 v2 而单独安装一个不匹配版本的 torchvision;版本不匹配可能导致导入错误或 C++ 算子不可用。

6.6.2.2 torchvision.transforms.v2 的使用

下面是图像分类训练集的变换管线。它将 PIL 的 RGB 图片随机裁剪到 224×224、随机水平翻转,然后转换成范围为 [0, 1] 的 float32 Tensor 并按 ImageNet 的均值/标准差归一化。最后输出形状为 (3, 224, 224) 的 Tensor:

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
# 60-数据增强torchvision.transforms.v2的使用.py
"""torchvision.transforms.v2 的图像分类数据增强示例。

本文件分别定义训练集与验证/测试集的变换管线:
训练集采用随机裁剪、随机翻转来增加样本多样性;
验证/测试集只采用固定的缩放与中心裁剪,保证每次评估同一张图片时结果一致。
"""

import torch # 提供 float32 数据类型,供 ToDtype 指定输出 Tensor 的类型。
from PIL import Image # Image.open() 用于从磁盘读取图片,返回 PIL.Image 对象。
from torchvision.transforms import v2 # torchvision 0.21 的新版图像变换 API。


# v2.Compose 会按照列表从上到下的顺序依次执行每一个变换。
# 训练集使用随机变换:同一张图片被多次读取时,裁剪区域和是否翻转都可能不同。
train_transform = v2.Compose([
# 将 PIL 图片转换为 torchvision 的图像 Tensor 表示,通道顺序变为 (C, H, W)。
# 此时像素通常仍是 uint8,数值范围仍是 [0, 255],尚未归一化。
v2.ToImage(), # 结果形状与数据类型:torch.Size([3, 1280, 1918]) torch.uint8

# 从原图随机选择一个区域,再把该区域缩放成 224×224。
# size 是最终输出的 (高度, 宽度),224×224 是许多 ImageNet 预训练分类模型的常见输入尺寸。
# scale=(0.8, 1.0) 表示随机裁取原图面积的 80%~100%,可模拟物体在画面中的不同大小和位置。
# antialias=True 在缩小图片时使用抗锯齿,减少缩放带来的锯齿与高频伪影。
v2.RandomResizedCrop(
size=(224, 224),
scale=(0.8, 1.0),
antialias=True,
), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.uint8

# 以 p=0.5(50%)的概率(剩余50%的图片保持原样)沿左右方向翻转图片(水平镜像翻转,就像照镜子一样)。
# 对“汽车类别”等左右翻转不改变标签含义的自然图像分类任务通常合理;
# 但文字识别、交通标志方向识别等任务不一定适用。
v2.RandomHorizontalFlip(p=0.5), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.uint8

# 把 uint8 的 [0, 255] 像素转换为 float32 的 [0, 1]。
# scale=True 很重要:它除以 255;若省略,转换 dtype 后数值仍会是 [0, 255]。
v2.ToDtype(torch.float32, scale=True), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.float32

# 对 RGB 三个通道分别进行标准化,计算公式为:
# output[channel] = (input[channel] - mean[channel]) / std[channel]。
# 下面的均值和标准差来自 ImageNet;使用 ImageNet 预训练权重微调时应保持一致。
# 具体做法:对图片的 RGB 三个通道分别做标准化,让输入数据的分布与 ImageNet 预训练模型训练时看到的数据分布一致。
# 新值 = (原值 - 该通道均值) / 该通道标准差
v2.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.float32
])


# 验证/测试集的目标是稳定、公平地评估模型,不能使用会改变输出的随机变换。
# 同一张图片每次经过 eval_transform 后都会得到相同结果。
eval_transform = v2.Compose([
# 同训练集:先把 PIL 图片转换为通道顺序为 (C, H, W) 的图像 Tensor。
v2.ToImage(), # 结果形状与数据类型:torch.Size([3, 1280, 1918]) torch.uint8

# 将图片的高度和宽度都缩放为 256。随后统一中心裁剪到 224×224,
# 是 ImageNet 分类评估中常用的预处理方式。
v2.Resize(size=(256, 256), antialias=True), # 结果形状与数据类型:torch.Size([3, 256, 256]) torch.uint8

# 从已经缩放到 256×256 的图片中心取出 224×224 区域。
# 它不含随机性,因此可使验证/测试结果可复现。
v2.CenterCrop(size=(224, 224)), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.uint8

# 与训练集相同:先转为 [0, 1] 范围的 float32,才能正确使用后续 Normalize。
v2.ToDtype(torch.float32, scale=True), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.float32

# 与训练集使用相同的 ImageNet 标准化参数,保证模型接收到的数据分布一致。
v2.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型:torch.Size([3, 224, 224]) torch.float32
])


# Image.open() 从当前工作目录读取训练图片,初始结果可能是 RGB、RGBA、灰度等不同模式。
# convert("RGB") 统一转换为 3 通道 RGB,避免灰度或带透明通道的图片与模型输入不匹配。
# 此处 ./Audi奥迪-0cdf5b5d0ce1_01.jpg 的颜色模型本来就是RGB,此图片来自如下网址
# https://www.kaggle.com/competitions/carvana-image-masking-challenge/data?select=train_hq.zip
image = Image.open("./Audi奥迪-0cdf5b5d0ce1_01.jpg").convert("RGB")

# 将训练图片依次送入 train_transform 中的五类变换,得到可直接输入分类模型的 Tensor。
# 因为包含随机裁剪和随机翻转,重复运行本行时 train_tensor 的具体像素通常会不同。
train_tensor = train_transform(image)
# print(f"train_transform: {train_transform}")
# print(f"train_tensor: {train_tensor}")

# 输出的形状是 (C, H, W) = (3, 224, 224),而非 PIL/NumPy 常见的 (H, W, C)。
# 数据类型是 float32;经过 Normalize 后,数值不再限制于 [0, 1]。
print(train_tensor.shape, train_tensor.dtype) # train_transform中内容全做后:torch.Size([3, 224, 224]) torch.float32


# 验证/测试可以使用另一张未参与训练的图片。它只经过确定性预处理,不进行随机增强。
image_test = Image.open("./Audi奥迪-0cdf5b5d0ce1_02.jpg").convert("RGB")

# 应用 eval_transform;对相同 image_test 重复执行时,eval_tensor 的值保持一致。
eval_tensor = eval_transform(image_test)
# print(f"\neval_transform: {eval_transform}")
# print(f"eval_tensor: {eval_tensor}")

# 验证/测试输出的形状和 dtype 与训练输入相同,可以直接喂给同一个模型。
print(eval_tensor.shape, eval_tensor.dtype) # eval_transform中内容全做后:torch.Size([3, 224, 224]) torch.float32

若使用 torchvision.datasets.ImageFolder,将训练变换传给 transform 参数即可;DataLoader 每次取样时会自动调用它:

1
2
3
4
5
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

train_dataset = ImageFolder("data/train", transform=train_transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

对于语义分割、目标检测等任务,不能只变换图像而不变换标注。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
2
3
4
5
6
jsy@jmac ~ % conda run -n ai_py312 python -m pip install -U albumentations

jsy@jmac ~ % conda run -n ai_py312 python -c "import cv2; import albumentations as A; print('OpenCV 版本:', cv2.__version__); print('Albumentations 版本:', A.__version__)"
# 输出如下:
OpenCV 版本: 5.0.0
Albumentations 版本: 2.0.8

Albumentations 会安装其需要的依赖。若运行图像读写或特定 OpenCV 变换时提示缺少 cv2,可按实际环境安装 opencv-python-headless(服务器、无 GUI 环境)或 opencv-python(桌面环境);两者通常只保留一个:

1
jsy@jmac ~ % conda run -n ai_py312 python -m pip install -U opencv-python-headless protobuf

6.6.3.2 Albumentations 的使用

6.6.3.2.1 Albumentations 的 Compose 简单操作

下面仍以图像分类为例。Albumentations 的 Compose 返回一个字典,图片位于 result["image"] 中。ToTensorV2 必须放在变换序列最后,因为前面的 Albumentations 变换处理的是 NumPy 图像:

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
# 61-数据增强albumentations的使用.py
"""使用 Albumentations 进行图像分类数据增强的完整示例。

本文件分别定义训练集与验证/测试集的变换管线:
训练集使用随机裁剪和水平翻转,增加训练样本的多样性;
验证/测试集只使用固定的缩放、中心裁剪与归一化,使同一张图片每次评估的结果一致。

为什么此文件没有像 torchvision.transforms.v2 一样在管线开始处把图片转为 Tensor?

torchvision.transforms.v2 是 PyTorch 生态的一部分,可直接处理 PIL.Image 或 PyTorch
Tensor,因此常见数据流是:PIL 图片 → v2.ToImage() → Tensor → 后续变换 → Tensor。

Albumentations 的大多数变换基于 NumPy / OpenCV,标准输入是形状为 (H, W, C) 的
numpy.ndarray,而不是 PyTorch Tensor。因此本文件采用:PIL 图片 → np.asarray()
→ HWC NumPy 数组 → Albumentations 几何/颜色变换 → ToTensorV2() → CHW Tensor。
ToTensorV2() 必须位于变换管线最后,之后的结果才能直接送入 PyTorch 模型。
"""

import albumentations as A # Albumentations 的图像变换算子,例如裁剪、翻转和归一化。
import numpy as np # np.asarray() 将 PIL 图片转换为 Albumentations 需要的 NumPy 数组。
from albumentations.pytorch import ToTensorV2 # 将 HWC NumPy 图片转换为 CHW PyTorch Tensor。
from PIL import Image # Image.open() 从磁盘读取图片,返回 PIL.Image 对象。


# A.Compose 按列表从上到下依次应用各变换。
# 本训练管线与 60-数据增强torchvision.transforms.v2的使用.py 的训练管线保持一致:
# 随机裁剪并缩放 → 随机水平翻转 → 转为 [0, 1] 的 float32 并按 ImageNet 参数标准化
# → 转换为 PyTorch Tensor。它不额外使用 ColorJitter 等颜色扰动。
# 因此,同一张图片多次读取时,裁剪区域和翻转状态可能不同,但颜色不会被随机改变。
train_transform = A.Compose([
# 随机从原图中裁取一个区域,再缩放成 224×224。
# size=(224, 224) 是输出图片的 (高度, 宽度),适合许多 ImageNet 预训练分类模型。
# scale=(0.8, 1.0) 表示裁取原图面积的 80%~100%,以模拟物体不同的大小与位置。
# p=1.0 表示每张训练图片都执行这一步。
A.RandomResizedCrop(
size=(224, 224),
scale=(0.8, 1.0),
# ratio 是随机裁剪区域的宽高比范围。
# (3/4, 4/3) 与 torchvision v2.RandomResizedCrop 的默认值一致,
# 允许裁出的区域略偏窄或略偏宽,但不会严重拉长、压扁物体。
# 这里限制的是“从原图随机选出的裁剪区域”的宽高比;该区域随后仍会缩放为
# 固定的 224×224 输出,因此最终输出图片的宽高比始终是 1:1。
ratio=(3 / 4, 4 / 3),
p=1.0,
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 以 50% 概率进行水平镜像翻转;另外 50% 概率保持原图不变。
# 对汽车分类等左右翻转不改变类别含义的任务通常适用;
# 对文字识别、车道方向判断等任务不一定适用。
A.HorizontalFlip(p=0.5), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 对 RGB 三个通道分别标准化,公式为:
# output[channel] = (input[channel] / 255.0 - mean[channel]) / std[channel]。
# A.Normalize 默认 max_pixel_value=255,因此会先将 uint8 的 [0, 255] 缩放为 [0, 1]。
# 它在这里同时完成了 torchvision v2 中以下两个连续步骤的等价工作:
# 1. v2.ToDtype(torch.float32, scale=True):uint8 [0, 255] -> float32 [0, 1];
# 2. v2.Normalize(mean=..., std=...):对三个 RGB 通道分别做标准化。
# 下面的 ImageNet 均值和标准差适用于使用 ImageNet 预训练权重进行微调的模型。
A.Normalize(
max_pixel_value=255.0, # 这行可省略,默认值就是 255.0;写上更清楚。
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.float32

# 必须放在最后:此前各变换处理 HWC NumPy 图片,
# 此处才转换为 PyTorch Tensor,并把形状由 (H, W, C) 改为 (C, H, W)。
ToTensorV2(), # 结果形状与数据类型: torch.Size([3, 224, 224]) torch.float32
])


# 验证/测试的目标是稳定、公平地评估模型,不能使用随机裁剪、翻转或颜色扰动。
# 这里每一步都是确定性的:同一张输入图片每次都会得到相同的输出 Tensor。
eval_transform = A.Compose([
# 固定把图片缩放到 256×256。这里的 height、width 是输出的高度与宽度。
# 此步骤没有随机性,即同一张输入图片、相同的参数、相同的插值设置下,每次执行都会得到相同的 256×256 输出;
# 每个输出像素对应原图的哪个位置,也是固定的。
# 后续再从中心裁剪为 224×224,是 ImageNet 分类评估中常见的预处理方式。
A.Resize(height=256, width=256), # 结果形状与数据类型: torch.Size([256, 256, 3]) torch.uint8

# 从 256×256 图片的正中心取出 224×224 区域。
# 这一步没有随机性,因此不会让验证/测试指标因图像变换而波动。
A.CenterCrop(height=224, width=224), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 与训练集使用相同的 ImageNet 标准化参数。
# 训练和评估的输入分布一致,模型才能正确使用训练期间学到的参数。
A.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.float32

# 最后将 HWC NumPy 图片变为 CHW 的 float32 PyTorch Tensor。
ToTensorV2(), # 结果形状与数据类型: torch.Size([3, 224, 224]) torch.float32
])


# ---------- 训练图片:随机数据增强 ----------

# Image.open() 读取图片后可能得到 RGB、RGBA、灰度等不同模式。
# convert("RGB") 统一为 3 通道 RGB,避免灰度图或透明通道图片与模型输入不匹配。
# np.asarray() 将 PIL 图片转为形状为 (H, W, C) 的 NumPy 数组,这是 Albumentations 的标准输入格式。
# 这一步与“60-数据增强torchvision.transforms.v2的使用.py”文件中 torchvision v2 示例的 v2.ToImage() 目的相似,都是转换图片表示;
# 但前者转换为 HWC NumPy 数组,后者转换为 CHW PyTorch Tensor。
image_np = np.asarray(Image.open("./Audi奥迪-0cdf5b5d0ce1_01.jpg").convert("RGB"))
print(f"type(image_np): {type(image_np)}") # <class 'numpy.ndarray'>
print(f"image_np shape: {image_np.shape}") # (1280, 1918, 3)
print(f"image_np dtype: {image_np.dtype}") # uint8
print(f"image_np: {image_np}") # [[[ 0 0 0] ... [ 0 0 0]] ... [[ 0 0 0] ... [ 0 0 0]]]

# A.Compose 的返回值是字典;传入 image=... 后,增强后的图片保存在 result["image"] 中。
# 由于包含随机变换,重复运行本行时 image_tensor 的具体像素一般会不同。
train_result = train_transform(image=image_np)
train_tensor = train_result["image"]

# 最终训练输入的形状为 (C, H, W) = (3, 224, 224),dtype 为 float32。
# 注意:经过 Normalize 后,数值不再限制在 [0, 1],出现负数是正常的。
print(train_tensor.shape, train_tensor.dtype) # torch.Size([3, 224, 224]) torch.float32


# ---------- 验证/测试图片:确定性预处理 ----------

# 读取另一张未参与训练的图片,并同样转换为 RGB 的 HWC NumPy 数组。
image_test_np = np.asarray(
Image.open("./Audi奥迪-0cdf5b5d0ce1_02.jpg").convert("RGB")
)

# 应用 eval_transform。这里没有随机算子,因此相同 image_test_np 每次输出都相同。
eval_result = eval_transform(image=image_test_np)
eval_tensor = eval_result["image"]

# 验证/测试输出与训练输入的形状和 dtype 一致,可直接输入同一个分类模型。
print(eval_tensor.shape, eval_tensor.dtype) # torch.Size([3, 224, 224]) torch.float32

6.6.3.2.2 分割任务时的单张变换

分割任务应在 Compose 调用时同时传入掩码,保证随机几何变换同步:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
segmentation_transform = A.Compose([
A.RandomResizedCrop(size=(224, 224),
scale=(0.8, 1.0),
ratio=(3 / 4, 4 / 3),
p=1.0,
),
A.HorizontalFlip(p=0.5),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])

# image_np 形状为 (H, W, 3),mask_np 形状为 (H, W)。
result = segmentation_transform(image=image_np, mask=mask_np)
image_tensor = result["image"] # 形状为 (3, 256, 256)。
mask_tensor = result["mask"] # 几何位置与 image_tensor 对齐;训练前再按损失函数要求调整 dtype。

如下给出一个二分类分割任务的相对完整可执行示例:

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
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
# 62-数据增强albumentations的使用-二分类分割同步变换.py
"""使用 Albumentations 对图像分割任务进行同步数据增强。

图像分类只有“整张图属于哪一类”的标签;语义分割则要求图片中每一个像素都有类别标签,
因此每张训练图片都必须有一张与它逐像素对齐的掩码(mask)。

本示例的关键是把 image 和 mask 一起传入同一个 A.Compose:

result = train_transform(image=image_np, mask=mask_np)

随机裁剪、缩放、水平翻转等几何变换会对 image 与 mask 使用完全相同的随机参数,
所以汽车在图片中的位置改变后,它在掩码中的位置也会同步改变。绝不能只增强图片而
不增强掩码,否则图片与标签会错位,模型无法学到正确的像素分类关系。

本例演示二分类分割:mask 中 0 表示背景、255 表示目标(例如汽车)。训练时会将
0/255 掩码转换为 0.0/1.0 的 float32 Tensor,适合与单通道输出及 BCE 类损失配合。
多类别分割时,应让 mask 像素直接保存类别编号 0、1、2……,并在训练前转为 long,
不要再除以 255。

哪些变换会处理 image,哪些会处理 mask?

Albumentations 为每个变换预先定义了可处理的目标类型(target)。当 Compose 接收到
特殊参数 image=...、mask=... 时,它会按每个变换的 target 定义自动路由数据,而不是
由用户手动为 image、mask 分别调用每一个变换。本文件各操作的对应关系如下:

+-------------------+------------+-----------+----------------------------+
| 操作 | 处理 image | 处理 mask | 原因 |
+-------------------+------------+-----------+----------------------------+
| RandomResizedCrop | 是 | 是 | 裁剪、缩放后的空间位置同步 |
| HorizontalFlip | 是 | 是 | 翻转方向同步 |
| Resize | 是 | 是 | 图片与掩码尺寸同步 |
| CenterCrop | 是 | 是 | 裁剪区域同步 |
| Normalize | 是 | 否 | 不能修改 mask 的类别编号 |
| ToTensorV2 | 是 | 是 | 转换为 PyTorch Tensor |
+-------------------+------------+-----------+----------------------------+

几何变换同步处理 image 与 mask 时,图片通常可使用双线性等插值来获得平滑结果;
mask 则默认使用最近邻插值,从而保持 0、255 等离散类别值,不生成 127 之类的中间标签。

如果一个样本有两个掩码,例如器官掩码和病灶掩码,可通过
A.Compose([...], additional_targets={"mask2": "mask"}) 声明第二张掩码,
之后调用 transform(image=image_np, mask=mask1_np, mask2=mask2_np)。两个掩码都会与
image 同步经历几何变换。
"""

from pathlib import Path # 用于在读取前检查图像和掩码文件是否存在。

import albumentations as A # 提供能同时处理图像与掩码的几何变换算子。
import numpy as np # np.asarray() 将 PIL 图片转换为 Albumentations 常用的 NumPy 数组。
import torch # 将二值掩码转换为适合二分类分割损失函数的 float32 Tensor。
from albumentations.pytorch import ToTensorV2 # 将 NumPy 图像/掩码转换为 PyTorch Tensor。
from PIL import Image # Image.open() 从磁盘读取图片或单通道掩码。


# A.Compose 按列表从上到下依次执行变换。
# 对 Compose 传入 image=...、mask=... 时,Albumentations 会让下面的几何变换同步作用于两者。
# 本训练管线与第 61 个分类示例的图像变换基本一致,但此处额外维护 mask 的空间对齐关系。
train_transform = A.Compose([
# 随机从原图和掩码的相同位置裁取一个区域,再同步缩放成 224×224。
# size=(224, 224) 是输出图片的 (高度, 宽度),适合许多 ImageNet 预训练分类模型。
# scale=(0.8, 1.0) 表示裁取原图面积的 80%~100%,以模拟物体不同的大小与位置。
# p=1.0 表示每张训练图片都执行这一步。
A.RandomResizedCrop(
size=(224, 224),
scale=(0.8, 1.0),
# ratio 是随机裁剪区域的宽高比范围。
# (3/4, 4/3) 与 torchvision v2.RandomResizedCrop 的默认值一致,
# 允许裁出的区域略偏窄或略偏宽,但不会严重拉长、压扁物体。
# 这里限制的是“从原图随机选出的裁剪区域”的宽高比;该区域随后仍会缩放为
# 固定的 224×224 输出,因此最终输出图片的宽高比始终是 1:1。
ratio=(3 / 4, 4 / 3),
p=1.0,
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 以 50% 概率将图片和掩码一起做水平镜像翻转;另外 50% 概率保持两者原样。
# 对汽车分类等左右翻转不改变类别含义的任务通常适用;
# 对文字识别、车道方向判断等任务不一定适用。
A.HorizontalFlip(p=0.5), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# A.Normalize 只对 image 的 RGB 三个通道做标准化,绝不会修改 mask 中的类别编号。
# 对 image 的计算公式为:
# output[channel] = (input[channel] / 255.0 - mean[channel]) / std[channel]。
# A.Normalize 默认 max_pixel_value=255,因此会先将 uint8 的 [0, 255] 缩放为 [0, 1]。
# 它在这里同时完成了 torchvision v2 中以下两个连续步骤的等价工作:
# 1. v2.ToDtype(torch.float32, scale=True):uint8 [0, 255] -> float32 [0, 1];
# 2. v2.Normalize(mean=..., std=...):对三个 RGB 通道分别做标准化。
# 下面的 ImageNet 均值和标准差适用于使用 ImageNet 预训练权重进行微调的模型。
A.Normalize(
max_pixel_value=255.0, # 这行可省略,默认值就是 255.0;写上更清楚。
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.float32

# 必须放在最后:此前各变换处理 HWC NumPy 图片,
# 此处将 image 由 (H, W, C) 转为 (C, H, W) 的 float32 Tensor;
# 同时将单通道 mask 由 (H, W) 转为 (H, W) 的 uint8 Tensor,不会对 mask 做归一化。
ToTensorV2(),
])


# 验证/测试的目标是稳定、公平地评估模型,不能使用随机裁剪、翻转或颜色扰动。
# 当 image 和 mask 一起传入时,下面的 Resize、CenterCrop 仍会同步处理两者;
# 每一步都是确定性的:同一组输入每次都会得到相同的 image Tensor 和 mask Tensor。
eval_transform = A.Compose([
# 固定把图片和掩码都缩放到 256×256。这里的 height、width 是输出的高度与宽度。
# 此步骤没有随机性,即同一张输入图片、相同的参数、相同的插值设置下,每次执行都会得到相同的 256×256 输出;
# 每个输出像素对应原图的哪个位置,也是固定的。
# 后续再从中心裁剪为 224×224,是 ImageNet 分类评估中常见的预处理方式。
A.Resize(height=256, width=256), # 结果形状与数据类型: torch.Size([256, 256, 3]) torch.uint8

# 从 256×256 图片和掩码的正中心同步取出 224×224 区域。
# 这一步没有随机性,因此不会让验证/测试指标因图像变换而波动。
A.CenterCrop(height=224, width=224), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 与训练集使用相同的 ImageNet 标准化参数;它仍只作用于 image,不会修改 mask。
# 训练和评估的输入分布一致,模型才能正确使用训练期间学到的参数。
A.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.float32

# 最后同步转换:image 变为 (3, 224, 224) float32 Tensor,mask 保持 (224, 224) uint8 Tensor。
ToTensorV2(),
])


# ---------- 需要准备的图像与掩码文件 ----------
#
# 每一张 RGB 图片必须有一张同名、同尺寸的单通道掩码:
# 训练图片:Audi奥迪-0cdf5b5d0ce1_01.jpg 训练掩码:Audi奥迪-0cdf5b5d0ce1_01_mask.png
# 验证图片:Audi奥迪-0cdf5b5d0ce1_02.jpg 验证掩码:Audi奥迪-0cdf5b5d0ce1_02_mask.png
# 来自: https://www.kaggle.com/competitions/carvana-image-masking-challenge/data?select=train_hq.zip
# https://www.kaggle.com/competitions/carvana-image-masking-challenge/data?select=train_masks.zip
# 本例为“汽车 / 非汽车”的二分类分割,掩码应保存为无损 PNG:
# - 黑色像素值 0:背景;
# - 白色像素值 255:目标汽车;
# - 掩码的高度、宽度必须与对应 RGB 图片完全相同;
# - 掩码不是原图的缩略图、轮廓图或涂了颜色的可视化图,而是每个像素都表示类别的标签图。
#
# PNG 不会像 JPEG 一样产生压缩伪影,所以适合保存离散的类别编号。不要把 mask 存为 JPG。
TRAIN_IMAGE_PATH = Path("./Audi奥迪-0cdf5b5d0ce1_01.jpg")
TRAIN_MASK_PATH = Path("./Audi奥迪-0cdf5b5d0ce1_01_mask.png") # 由原“0cdf5b5d0ce1_01_mask.gif”文件导出或另存为得到
EVAL_IMAGE_PATH = Path("./Audi奥迪-0cdf5b5d0ce1_02.jpg")
EVAL_MASK_PATH = Path("./Audi奥迪-0cdf5b5d0ce1_02_mask.png")


def load_image_and_mask(image_path, mask_path):
"""读取一对 RGB 图片与二值掩码,并检查它们能否逐像素对应。"""
# 先检查文件存在性。提示会明确告诉使用者缺少的是图片还是掩码。
for path in (image_path, mask_path):
if not path.is_file():
raise FileNotFoundError(
f"找不到文件:{path}\n"
"请在本脚本同目录准备 RGB 图片及同尺寸的 0/255 PNG 掩码,"
"或按实际文件名修改上方的 *_IMAGE_PATH、*_MASK_PATH。"
)

# 图片统一转为 RGB,得到形状为 (H, W, 3) 的 uint8 NumPy 数组。
# 这是 Albumentations 对普通彩色图像的标准输入格式。
image_np = np.asarray(Image.open(image_path).convert("RGB"))

# 掩码转为 L(灰度)模式,得到形状为 (H, W) 的单通道 uint8 NumPy 数组。
# 不要对 mask 使用 convert("RGB"),否则会得到三个重复通道,不符合本例的二值标签格式。
mask_np = np.asarray(Image.open(mask_path).convert("L"))

# 只有高度、宽度完全相同,mask 的每个像素才表示 image 同一位置的类别。
if image_np.shape[:2] != mask_np.shape:
raise ValueError(
f"图片与掩码尺寸不一致:图片为 {image_np.shape[:2]},"
f"掩码为 {mask_np.shape}。请使用与原图同尺寸的掩码。"
)

# 二值掩码只允许 0 与 255;提前检查可以避免把彩色可视化图误当作标签图。
# np.unique 会把输入数组展平(Flatten),提取所有不重复的元素,并自动按从小到大的顺序排列
unique_values = np.unique(mask_np)
if not np.isin(unique_values, [0, 255]).all():
raise ValueError(
f"本示例要求二值掩码像素只能是 0 或 255,但读到的值包括:{unique_values[:10]}。"
)

return image_np, mask_np


def binary_mask_to_target(mask_tensor):
"""将 ToTensorV2 输出的 0/255 uint8 掩码转换为 BCE 训练目标。"""
# ToTensorV2 对单通道 mask 的输出形状为 (H, W),dtype 通常是 uint8。
# unsqueeze(0) 增加通道维度,得到 (1, H, W),与单通道分割模型的输出形状匹配。
# 再除以 255,使背景从 0 变为 0.0,汽车从 255 变为 1.0。
return mask_tensor.unsqueeze(0).to(torch.float32) / 255.0


def show_transformed_pair(split_name, transform, image_path, mask_path):
"""读取一对样本、同步增强,并打印增强后图片与掩码的形状。"""
image_np, mask_np = load_image_and_mask(image_path, mask_path)

# 关键调用:必须同时传入 image 和 mask。
# 对 train_transform 而言,随机裁剪和随机翻转的参数会同时用于这两个输入;
# 对 eval_transform 而言,缩放和中心裁剪同样会同步应用。
result = transform(image=image_np, mask=mask_np)

# A.Compose 返回字典:result["image"] 是经过 Normalize + ToTensorV2 的图片,
# result["mask"] 是仅经过几何变换 + ToTensorV2 的掩码,类别值 0/255 保持不变。
image_tensor = result["image"]
raw_mask_tensor = result["mask"]
target_mask_tensor = binary_mask_to_target(raw_mask_tensor)

print(f"\n{split_name}:")
print(f"原始图片形状 / dtype:{image_np.shape} / {image_np.dtype}")
print(f"原始掩码形状 / dtype:{mask_np.shape} / {mask_np.dtype}")
print(f"增强后图片形状 / dtype:{image_tensor.shape} / {image_tensor.dtype}")
print(f"增强后原始掩码形状 / dtype:{raw_mask_tensor.shape} / {raw_mask_tensor.dtype}")
print(
"BCE 目标掩码形状 / dtype / 像素值:"
f"{target_mask_tensor.shape} / {target_mask_tensor.dtype} / "
f"{torch.unique(target_mask_tensor).tolist()}"
)

return image_tensor, target_mask_tensor


def main():
# 训练集:每次运行都可能使用不同随机裁剪/翻转,但图片与掩码始终空间对齐。
show_transformed_pair(
"训练样本(随机同步增强)",
train_transform,
TRAIN_IMAGE_PATH,
TRAIN_MASK_PATH,
)

# 验证/测试集:没有随机变换,重复运行时得到相同的图片与掩码结果。
show_transformed_pair(
"验证/测试样本(确定性同步预处理)",
eval_transform,
EVAL_IMAGE_PATH,
EVAL_MASK_PATH,
)


if __name__ == "__main__":
main()


"""
### 运行结果示例(每次运行训练样本可能不同,但验证/测试样本固定):
(ai_py312) jsy@jmac 06第六章 % python 62-数据增强albumentations的使用-二分类分割同步变换.py
/opt/anaconda3/envs/ai_py312/lib/python3.12/site-packages/albumentations/check_version.py:147: UserWarning: Error fetching version info <urlopen error _ssl.c:993: The handshake operation timed out>
data = fetch_version_info()

训练样本(随机同步增强):
原始图片形状 / dtype:(1280, 1918, 3) / uint8
原始掩码形状 / dtype:(1280, 1918) / uint8
增强后图片形状 / dtype:torch.Size([3, 224, 224]) / torch.float32
增强后原始掩码形状 / dtype:torch.Size([224, 224]) / torch.uint8
BCE 目标掩码形状 / dtype / 像素值:torch.Size([1, 224, 224]) / torch.float32 / [0.0, 1.0]

验证/测试样本(确定性同步预处理):
原始图片形状 / dtype:(1280, 1918, 3) / uint8
原始掩码形状 / dtype:(1280, 1918) / uint8
增强后图片形状 / dtype:torch.Size([3, 224, 224]) / torch.float32
增强后原始掩码形状 / dtype:torch.Size([224, 224]) / torch.uint8
BCE 目标掩码形状 / dtype / 像素值:torch.Size([1, 224, 224]) / torch.float32 / [0.0, 1.0]
"""
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
2
3
4
5
6
7
8
9
10
11
12
13
data/
├── train/
│ ├── images/
│ │ ├── 0001.jpg
│ │ └── 0002.jpg
│ └── masks/
│ ├── 0001_mask.png
│ └── 0002_mask.png
└── eval/
├── images/
│ └── 0003.jpg
└── masks/
└── 0003_mask.png

每张掩码必须与对应 RGB 图片宽高完全一致。本节为“目标 / 背景”的二分类分割示例:mask 使用单通道无损 PNG,背景像素为 0,目标像素为 255。不要用 JPG 保存 mask,因为有损压缩可能产生 1、254 等不属于类别的值。

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
# 63-数据增强albumentations的使用-二分类分割批量加载.py
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset

import albumentations as A
from albumentations.pytorch import ToTensorV2

# A.Compose 按列表从上到下依次执行变换。
# 对 Compose 传入 image=...、mask=... 时,Albumentations 会让下面的几何变换同步作用于两者。
# 本训练管线与第 61 个分类示例的图像变换基本一致,但此处额外维护 mask 的空间对齐关系。
train_transform = A.Compose([
# 随机从原图和掩码的相同位置裁取一个区域,再同步缩放成 224×224。
# size=(224, 224) 是输出图片的 (高度, 宽度),适合许多 ImageNet 预训练分类模型。
# scale=(0.8, 1.0) 表示裁取原图面积的 80%~100%,以模拟物体不同的大小与位置。
# p=1.0 表示每张训练图片都执行这一步。
A.RandomResizedCrop(
size=(224, 224),
scale=(0.8, 1.0),
# ratio 是随机裁剪区域的宽高比范围。
# (3/4, 4/3) 与 torchvision v2.RandomResizedCrop 的默认值一致,
# 允许裁出的区域略偏窄或略偏宽,但不会严重拉长、压扁物体。
# 这里限制的是“从原图随机选出的裁剪区域”的宽高比;该区域随后仍会缩放为
# 固定的 224×224 输出,因此最终输出图片的宽高比始终是 1:1。
ratio=(3 / 4, 4 / 3),
p=1.0,
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 以 50% 概率将图片和掩码一起做水平镜像翻转;另外 50% 概率保持两者原样。
# 对汽车分类等左右翻转不改变类别含义的任务通常适用;
# 对文字识别、车道方向判断等任务不一定适用。
A.HorizontalFlip(p=0.5), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# A.Normalize 只对 image 的 RGB 三个通道做标准化,绝不会修改 mask 中的类别编号。
# 对 image 的计算公式为:
# output[channel] = (input[channel] / 255.0 - mean[channel]) / std[channel]。
# A.Normalize 默认 max_pixel_value=255,因此会先将 uint8 的 [0, 255] 缩放为 [0, 1]。
# 它在这里同时完成了 torchvision v2 中以下两个连续步骤的等价工作:
# 1. v2.ToDtype(torch.float32, scale=True):uint8 [0, 255] -> float32 [0, 1];
# 2. v2.Normalize(mean=..., std=...):对三个 RGB 通道分别做标准化。
# 下面的 ImageNet 均值和标准差适用于使用 ImageNet 预训练权重进行微调的模型。
A.Normalize(
max_pixel_value=255.0, # 这行可省略,默认值就是 255.0;写上更清楚。
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.float32

# 必须放在最后:此前各变换处理 HWC NumPy 图片,
# 此处将 image 由 (H, W, C) 转为 (C, H, W) 的 float32 Tensor;
# 同时将单通道 mask 由 (H, W) 转为 (H, W) 的 uint8 Tensor,不会对 mask 做归一化。
ToTensorV2(),
])


# 验证/测试的目标是稳定、公平地评估模型,不能使用随机裁剪、翻转或颜色扰动。
# 当 image 和 mask 一起传入时,下面的 Resize、CenterCrop 仍会同步处理两者;
# 每一步都是确定性的:同一组输入每次都会得到相同的 image Tensor 和 mask Tensor。
eval_transform = A.Compose([
# 固定把图片和掩码都缩放到 256×256。这里的 height、width 是输出的高度与宽度。
# 此步骤没有随机性,即同一张输入图片、相同的参数、相同的插值设置下,每次执行都会得到相同的 256×256 输出;
# 每个输出像素对应原图的哪个位置,也是固定的。
# 后续再从中心裁剪为 224×224,是 ImageNet 分类评估中常见的预处理方式。
A.Resize(height=256, width=256), # 结果形状与数据类型: torch.Size([256, 256, 3]) torch.uint8

# 从 256×256 图片和掩码的正中心同步取出 224×224 区域。
# 这一步没有随机性,因此不会让验证/测试指标因图像变换而波动。
A.CenterCrop(height=224, width=224), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.uint8

# 与训练集使用相同的 ImageNet 标准化参数;它仍只作用于 image,不会修改 mask。
# 训练和评估的输入分布一致,模型才能正确使用训练期间学到的参数。
A.Normalize(
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
), # 结果形状与数据类型: torch.Size([224, 224, 3]) torch.float32

# 最后同步转换:image 变为 (3, 224, 224) float32 Tensor,mask 保持 (224, 224) uint8 Tensor。
ToTensorV2(),
])

class BinarySegmentationDataset(Dataset):
"""读取一对 image/mask,并在 __getitem__ 中同步执行数据增强。"""

def __init__(self, image_dir, mask_dir, transform):
self.image_dir = Path(image_dir)
self.mask_dir = Path(mask_dir)
self.transform = transform

# 读取 images 目录中全部 JPG 图片,并按文件名排序,使样本顺序稳定。
self.image_paths = sorted(self.image_dir.glob("*.jpg"))
if not self.image_paths:
raise FileNotFoundError(f"没有在目录中找到 JPG 图片:{self.image_dir}")

# 例如 images/0001.jpg 对应 masks/0001_mask.gif。
self.mask_paths = [
self.mask_dir / f"{image_path.stem}_mask.gif"
for image_path in self.image_paths
]

# 提前检查每张图片是否都有对应掩码,避免训练中途才发现漏标注。
for mask_path in self.mask_paths:
if not mask_path.is_file():
raise FileNotFoundError(f"找不到图片对应的掩码:{mask_path}")

def __len__(self):
# DataLoader 据此确定一个 epoch 的样本数。
return len(self.image_paths)

def __getitem__(self, index):
image_path = self.image_paths[index]
mask_path = self.mask_paths[index]

# image_np 为 (H, W, 3) 的 RGB uint8 NumPy 数组。
image_np = np.asarray(Image.open(image_path).convert("RGB"))

# mask_np 为 (H, W) 的单通道 uint8 数组,像素值应为 0 或 255。
# mask 不能 convert("RGB"),否则会产生三个重复通道。
mask_np = np.asarray(Image.open(mask_path).convert("L"))

# 掩码的每个像素必须对应图片同一位置,故高、宽必须相同。
if image_np.shape[:2] != mask_np.shape:
raise ValueError(
f"图片与掩码尺寸不一致:{image_path.name} 为 {image_np.shape[:2]},"
f"{mask_path.name} 为 {mask_np.shape}"
)

# 关键步骤:同一次调用同时传入 image 与 mask。
# train_transform 的随机裁剪、翻转会使用同一套随机参数;
# eval_transform 的缩放、中心裁剪也会同步处理两者。
result = self.transform(image=image_np, mask=mask_np)
image_tensor = result["image"] # (3, 224, 224),float32,已 Normalize。
mask_tensor = result["mask"] # (224, 224),uint8,仍是 0/255 离散类别。

# 为 BCEWithLogitsLoss 准备二分类目标:
# (H, W) -> (1, H, W),并将 uint8 0/255 -> float32 0.0/1.0。
mask_tensor = mask_tensor.unsqueeze(0).to(torch.float32) / 255.0
return image_tensor, mask_tensor

# 本文件路径为 source/06第六章/66-...py,parents[1] 即 source/ 目录。
DATA_ROOT = Path(__file__).resolve().parents[1] / "carvana-image-masking-challenge-data"
# train_transform、eval_transform 是第 62 个文件定义的 Albumentations Compose 对象。
# 训练集打乱顺序;验证集不打乱,便于稳定评估。
train_dataset = BinarySegmentationDataset(
f"{DATA_ROOT}/train/images", f"{DATA_ROOT}/train/masks", transform=train_transform
)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)

eval_dataset = BinarySegmentationDataset(
f"{DATA_ROOT}/eval/images", f"{DATA_ROOT}/eval/masks", transform=eval_transform
)
eval_loader = DataLoader(eval_dataset, batch_size=8, shuffle=False)


# DataLoader 将多个已同步增强完成的样本组成 batch;最后一个 batch 可能小于 8。
for images, masks in train_loader:
print(images.shape) # torch.Size([8, 3, 224, 224])
print(masks.shape) # torch.Size([8, 1, 224, 224])
break


for images, masks in eval_loader:
print(images.shape) # torch.Size([8, 3, 224, 224])
print(masks.shape) # torch.Size([8, 1, 224, 224])
break

也就是说,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)

找到一组优秀的超参数通常是深度学习实验中最耗费时间的工作之一,常见的调优策略包括:

  1. 网格搜索(Grid Search):穷举法。在预先设定的参数组合列表里遍历测试所有可能(耗时极长)。
  2. 随机搜索(Random Search):在参数空间内随机抽样组合。通常比网格搜索更高效,能更快覆盖更大的搜索空间。
  3. 贝叶斯优化(Bayesian Optimization):根据已测试组合的结果,用概率模型预测下一组最有潜力的超参数组合(自动化调参库如 Optuna、Ray Tune 常用的核心策略)。
  4. 凭经验与手动调整(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
2
3
4
5
import argparse

parser = argparse.ArgumentParser(
description="训练汽车品牌分类模型。",
)

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
2
3
4
5
parser = argparse.ArgumentParser(
description="训练汽车品牌分类模型。",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
epilog="示例:python train.py --epochs 20 --lr 0.0003",
)

此时执行:

1
python train.py --help

不会开始训练,而是显示用法、description、后续通过 add_argument() 注册的参数规格。这是检查“参数名是否写对、默认值是什么”的首选方法。

6.7.3.3 ArgumentParser.add_argument常用参数

6.7.3.3.1 常见参数概述

以下是执行ArgumentParser.add_argument(xxx)添加参数规格时经常使用到的参数及其解释: 先看一个典型的训练超参数:

1
2
3
4
5
6
7
8
9
10
11
12
parser = argparse.ArgumentParser(
description="按品牌整理 Stanford Cars 数据集。",
epilog='示例: python 65-按品牌整理StanfordCars数据集.py',
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)

parser.add_argument(
"--lr",
type=float,
default=1e-3,
help="AdamW 优化器的初始学习率。",
)

它规定了用户可在命令行传入 --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
2
3
4
5
6
7
# 位置参数:必须写,且没有 -- 前缀。
# 命令示例:python train.py data/train
parser.add_argument("data_dir", type=Path, help="训练数据根目录。")

# 可选参数:可以省略;省略时采用 default。
# 命令示例:python train.py data/train --epochs 20
parser.add_argument("--epochs", type=int, default=10, help="训练轮数。")

布尔参数是初学者容易出错的地方。不要写 type=bool:命令行中的 "False" 是一个非空字符串,转换时仍可能得到 True。更清晰的写法是:

1
2
3
4
5
6
7
8
9
10
# 默认关闭;仅当命令行明确出现 --amp 时才开启。
parser.add_argument("--amp", action="store_true", help="启用 CUDA 自动混合精度。")

# 默认开启;仅当命令行明确出现 --no-augment 时才关闭。
parser.add_argument(
"--no-augment",
action="store_false",
dest="augment",
help="关闭训练数据增强。",
)

上例中,未传 --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
2
3
4
5
6
7
8
9
10
11
12
13
14
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--cuda', action='store_true', default=False, help='是否使用 CUDA')
parser.add_argument('--no-augment', action='store_false', dest='augment', default=True)

# 场景 A: 命令行不传任何参数 -> python script.py
args = parser.parse_args([])
print(args.cuda) # Output: False (未传参数,取 default)
print(args.augment) # Output: True (未传参数,取 default)

# 场景 B: 命令行传入 --cuda -> python script.py --cuda
args = parser.parse_args(['--cuda'])
print(args.cuda) # Output: True (触发 store_true,覆盖 default)

注意:如果不显式写 default,store_true 默认的 default 本身就是 False,store_false 默认的 default 本身就是 True。但显式指定 default 能让代码意图更清晰,或者在特殊需求下反向定义(例如将 action="store_true" 的 default 设为 None,以区分“用户未传”与“用户显式关闭/开启”)。

2. action="store"(默认动作,存入单一值)
  • 原理:store 是默认动作。如果命令行指定了该参数,则保存用户给定的值;如果命令行未指定,则保存 default 的值。
1
2
3
4
5
6
7
8
9
parser.add_argument('--lr', action='store', type=float, default=0.001)

# 命令行未指定 --lr
args = parser.parse_args([])
print(args.lr) # Output: 0.001

# 命令行指定 --lr 0.01
args = parser.parse_args(['--lr', '0.01'])
print(args.lr) # Output: 0.01
3. action="append"(追加到列表)
  • 原理:每次在命令行遇到该参数时,将值追加(append)到列表中。
  • 与 default 的冲突点/易错点:
    • 未指定参数:变量值为 default 的原始值(建议设为列表或 None)。
    • 指定参数:用户传入的值会追加在 default 列表的末尾,而不会覆盖 default 列表!
1
2
3
4
5
6
7
8
9
parser.add_argument('--seed', action='append', type=int, default=[42])

# 场景 A: 命令行未指定 -> 取 default 列表
args = parser.parse_args([])
print(args.seed) # Output: [42]

# 场景 B: 命令行指定 --seed 100 -> 追加到 default 列表之后!
args = parser.parse_args(['--seed', '100'])
print(args.seed) # Output: [42, 100] <-- 经常被误以为会变成 [100]

最佳实践:如果使用 action="append",通常建议设置 default=None 或在解析后再做判定,避免默认列表被污染。

4. action="count"(统计出现次数)
  • 原理:统计标志在命令行出现的次数(常用于日志详细程度 --verbose 或 -vvv)。
  • 行为:
    • 未传入时,变量值为 default。
    • 传入 \(N\) 次时,变量值为 default + N(如果 default 为整型)或直接计数(如果 default=None,首次出现从 1 开始计)。
1
2
3
4
5
6
7
8
9
10
11
12
parser.add_argument('-v', '--verbose', action='count', default=0)

# 未指定
args = parser.parse_args([])
print(args.verbose) # Output: 0

# 指定 -vv
args = parser.parse_args(['-vv'])
print(args.verbose) # Output: 2
# 指定 -vvv
args = parser.parse_args(['-vvv'])
print(args.verbose) # Output: 3
5. action="store_const"(存储常数)
  • 原理:当指定参数标志时,将 const 参数赋予的值存入;未指定时,存入 default 的值。
1
2
3
4
5
6
7
parser.add_argument('--mode', action='store_const', const='fast', default='slow')

args = parser.parse_args([])
print(args.mode) # Output: 'slow' (未指定,取 default)

args = parser.parse_args(['--mode'])
print(args.mode) # Output: 'fast' (指定了,取 const)
总结对照表
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
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
# 70-参数解析器经常性用法.py
from __future__ import annotations

import argparse
from pathlib import Path


def parse_args() -> argparse.Namespace:
"""定义训练脚本接受的超参数,并解析当前命令行。"""
parser = argparse.ArgumentParser(
description="训练一个汽车品牌图像分类模型。",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
epilog=(
"示例:python train.py --data-dir source/data "
"--epochs 20 --batch-size 32 --lr 0.0003 --amp"
),
)

# 数据相关参数:data_dir 没有合理的通用默认路径,所以设置为必传。
parser.add_argument(
"--data-dir",
type=Path,
required=True,
metavar="DIR",
help="包含 train/ 与 eval/ 子目录的数据根目录。",
)
parser.add_argument(
"--image-size",
type=int,
default=224,
help="输入模型前,图片缩放/裁剪后的边长。",
)

# 训练过程相关的超参数。
parser.add_argument("--epochs", type=int, default=10, help="训练轮数。")
parser.add_argument("--batch-size", type=int, default=32, help="每个 batch 的样本数。")
parser.add_argument("--lr", type=float, default=1e-3, help="优化器的初始学习率。")
parser.add_argument(
"--weight-decay",
type=float,
default=1e-4,
help="AdamW 的权重衰减系数。",
)
parser.add_argument(
"--optimizer",
choices=["sgd", "adamw"],
default="adamw",
help="优化器类型。",
)

# 运行方式相关的参数。
parser.add_argument(
"--device",
choices=["auto", "cpu", "cuda"],
default="auto",
help="训练设备;auto 表示由后续代码自动选择。",
)
parser.add_argument("--seed", type=int, default=42, help="随机种子。")
parser.add_argument("--amp", action="store_true", help="启用 CUDA 自动混合精度。")
parser.add_argument(
"--no-augment",
action="store_false",
dest="augment",
help="关闭训练数据增强。",
)

args = parser.parse_args()

# argparse 负责基本的“字符串 -> 类型”和 choices 检查;参数之间的业务规则
# 则适合在解析后自行检查,以便给出符合训练语境的提示。
if args.epochs <= 0:
parser.error("--epochs 必须是正整数。")
if args.batch_size <= 0:
parser.error("--batch-size 必须是正整数。")
if args.image_size <= 0:
parser.error("--image-size 必须是正整数。")
if args.lr <= 0:
parser.error("--lr 必须大于 0。")
if args.weight_decay < 0:
parser.error("--weight-decay 不能小于 0。")

return args


def main() -> None:
args = parse_args()

# parse_args() 返回 Namespace;可用点运算符读取其中每个参数。
print("数据目录:", args.data_dir)
print("训练轮数:", args.epochs)
print("batch 大小:", args.batch_size)
print("学习率:", args.lr)
print("优化器:", args.optimizer)
print("使用数据增强:", args.augment)
print("启用 AMP:", args.amp)

# 真实训练脚本通常会在这里继续:
# 1. 按 args.data_dir 创建 ImageFolder / Dataset 和 DataLoader;
# 2. 用 args.batch_size 创建 DataLoader;
# 3. 用 args.optimizer、args.lr、args.weight_decay 创建优化器;
# 4. 循环 args.epochs 轮,并在 args.amp 为 True 时启用 AMP。


if __name__ == "__main__":
main()

将代码保存为 70-参数解析器经常性用法.py 后,可按下面方式运行:

1
2
3
4
5
6
7
8
# 使用默认的训练超参数;--data-dir 是 required=True 的必传参数。
python 70-参数解析器经常性用法.py --data-dir source/data

# 覆盖部分默认值,并打开 AMP。
python 70-参数解析器经常性用法.py --data-dir source/data --epochs 20 --batch-size 16 --lr 0.0003 --amp

# 查看所有参数说明、可选值和默认值。
python 70-参数解析器经常性用法.py --help

第二条命令解析后的部分结果等价于:

1
2
3
4
5
args.epochs == 20
args.batch_size == 16
args.lr == 0.0003
args.amp is True
args.augment is True # 未传入 --no-augment,故保持默认开启

命令行参数只负责“本次运行的配置入口”。为了让日后能够复现一次实验,实际训练时还应将 vars(args)(或 args 的内容)连同模型权重、训练日志一起保存到 JSON、YAML 或 checkpoint 中。

6.7.3.5 总结

可以看到,argparse模块为命令行的传入与解析提供了方便、实用的API接口,极大地方便深度学习中超参数的传入与解析。

一般地,为了复现或继承训练结果,深度学习会将超参数以json或dict或yaml等格式文件进行保存,模型结构与权重参数以pth等格式文件进行保存。


PyTorch06-PyTorch进阶训练技巧
https://jiangsanyin.github.io/2026/09/21/PyTorch06-PyTorch进阶训练技巧/
作者
sanyinjiang
发布于
2026年9月21日
许可协议