模型训练 warm-up, consine schedule, 以及 Automatic Mixed Precision

在模型训练过程中，模型预热（Warm-up）和余弦退火调度（Consine schedule）对于提高模型的训练效果和稳定性非常有帮助，它们可以优化学习率的变化，进而影响模型的学习过程。自动混合精度训练（Automatic Mixed Precision, AMP），这是一种优化训练过程的技术，尤其在使用 NVIDIA GPU 时。这种方法可以提高计算速度，同时减少所需的内存带宽，这是通过在训练神经网络时灵活使用单精度（float32）和半精度（float16）来实现的

预热（Warm-up）

预热期是训练开始的几个迭代或者几个周期，在这个阶段，学习率从一个相对较低的值逐渐增加到一个预定的较高值。预热的目的包括：

• 防止模型早期发散：在神经网络训练初期，权重是随机初始化的，此时若直接使用较大的学习率，可能会导致权重更新过大，从而使得模型输出快速发散，尤其是在损失函数的梯度非常大的情况下。
• 逐步适应训练：通过慢慢增加学习率，模型可以逐渐适应数据集的特点和复杂性，这有助于提高模型训练的稳定性和效率。
  提高训练效果：在一些情况下，预热可以帮助模型达到更好的最终性能，尤其是在使用大型模型和复杂数据集时。

    if warmup_epochs > 0:
        warmup_schedule = np.linspace(0, lr, warmup_epochs)

余弦退火调度（Consine schedule）

余弦退火是一种调整学习率的方法，它模拟了退火过程中的冷却曲线，即学习率随训练进程先减小后平稳。具体来说：

• 避免局部最小：通过模拟“退火”过程，余弦退火允许学习率初期较高，以跳出可能的局部最小值，然后逐渐降低学习率以细化到全局最小或更好的局部最小。
• 优化收敛：随着训练的进行，模型的参数逐渐接近最优，此时需要较小的学习率来进行微调，避免过大的参数更新导致“越过”最佳解。
• 控制结束阶段的学习率：通过模拟余弦曲线，学习率在训练的最后阶段可以非常平稳，有助于模型在最优解附近稳定下来。

在实际使用中，可以使用pytorch框架内置函数:

import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = optim.Adam(model.parameters(), lr=0.01)  # 初始学习率
# 设置调度器
scheduler = CosineAnnealingLR(optimizer, T_max=10, eta_min=0.001)  # T_max 是周期，eta_min 是最小学习率

data_loader = [torch.randn(10, 10) for _ in range(100)] 
# 训练循环
for epoch in range(10):
    for inputs in data_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = outputs.sum()  # 使用简单的求和作为损失函数示例
        loss.backward()
        optimizer.step()
    scheduler.step()  # 更新学习率

    print(f"Epoch {epoch+1}: Current learning rate = {scheduler.get_last_lr()[0]}")

也可以自定义consine scheduler

$${\eta }_t={\eta }_{\min }+\frac{1}{2}\left({\eta }_{\max }-{\eta }_{\min }\right)\left(1+\cos \left(\frac{t\pi }{T}\right)\right)$$

def consine_scheduler(base_value, final_value, epochs, niter_per_ep):
	# final_value should be smaller than base_value
	iters = np.arange(epochs * niter_per_ep)
    schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))
    assert len(schedule) == epochs * niter_per_ep
    return schedule

也可以把预热和consine_scheduler放在一起：

def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0):
    warmup_schedule = np.array([])
    warmup_iters = warmup_epochs * niter_per_ep
    if warmup_epochs > 0:
        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)

    iters = np.arange(epochs * niter_per_ep - warmup_iters)
    schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))

    schedule = np.concatenate((warmup_schedule, schedule))
    assert len(schedule) == epochs * niter_per_ep
    return schedule

自动混合精度 (Automatic Mixed Precision)

混合精度训练结合了两种数据类型：

• 单精度（Float32）：传统上，深度学习模型使用单精度格式进行计算和存储参数。单精度提供了足够的数值范围和精度，适合保证大多数训练任务的稳定性和收敛性。
• 半精度（Float16）：半精度提供较低的数值精度，但在现代GPU（如NVIDIA的Volta及更高级的架构）上，使用半精度进行计算速度更快，且消耗更少的内存。

torch.cuda.amp.autocast 允许自动选择最适合操作的数据类型，它能在运行时动态地将数据类型从Float32转换为Float16，从而加速训练并减少内存使用，而不牺牲模型的性能或精度。这种转换对大多数操作是透明的。

with autocast(): # 默认 enabled=True
        outputs = model(inputs.cuda())
        loss = loss_fn(outputs, targets.cuda())

在 PyTorch 中，torch.cuda.amp.GradScaler 是配合自动混合精度 (autocast) 使用的工具，旨在处理半精度计算中可能出现的数值不稳定问题。由于半精度浮点数的动态范围有限，计算过程中的梯度可能出现太小而无法在半精度中精确表示的情况，这就是所谓的“梯度欠流”（underflow）和“梯度溢出”（overflow）。GradScaler 通过动态调整梯度的规模（即缩放因子）来帮助避免这些问题。具体来说，它会自动调整梯度的大小，以确保它们不会因为太小而在半精度表示中丢失，也不会因为太大而溢出。

import torch
from torch import nn
from torch.cuda.amp import autocast, GradScaler
import torch.optim as optim

# 定义模型和优化器
model = nn.Linear(10, 1).cuda()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 初始化GradScaler
scaler = GradScaler()

for data, target in data_loader:  # 假设data_loader已经定义
    optimizer.zero_grad()
    
    # 自动混合精度块
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # 放大梯度
    scaler.scale(loss).backward()
    
    # 优化器更新前，先尝试减小梯度规模
    scaler.step(optimizer)
    
    # 更新缩放器
    scaler.update()