3D模型训练发散的表现通常比较剧烈:训练损失在几十个迭代内从正常值突然攀升到NaN,或者权重范数在一个epoch内扩大几个数量级。与2D卷积网络不同,3D模型往往参数更多、单卡可容纳的批次更小、梯度噪声更明显,这些因素叠加后,不恰当的学习率与优化器组合很容易把训练推向发散。

调参时直接无脑减小学习率并不是最优解。例如学习率设置得过低虽然暂时不会发散,但训练会变得非常缓慢,甚至陷入尖锐极小值导致泛化变差。正确的做法是先通过损失曲线和梯度统计判断发散类型,再针对性地调整优化器参数与学习率调度策略。
一、先定位发散类型:权重爆炸还是更新震荡
训练发散可以粗略分为两类:一类是权重爆炸,表现为权重范数持续增大,最终损失变成NaN;另一类是更新震荡,损失在高低之间剧烈波动但不一定变为NaN,通常说明学习率处于临界值附近,优化器二阶矩估计不稳定。
区分这两类有助于决定后续动作。权重爆炸往往意味着单步参数更新量过大,或者梯度本身出现异常值。此时应优先考虑梯度裁剪和降低学习率;更新震荡则更多与优化器的自适应学习率、批次噪声以及学习率调度节奏有关,可以尝试调整Adam的epsilon、切换到SGD,或者延长预热周期。
在PyTorch中可以在训练循环中记录梯度范数和权重范数。梯度范数突然超过某个阈值,例如大于10或100,基本可以判断为该批次出现了异常梯度;权重范数如果每个step都以超过5%的速度增长,说明更新量已经超出稳定范围。下面代码展示了如何安全地获取这两项指标:
total_grad_norm = 0.0
for p in model.parameters():
if p.grad is not None:
param_grad_norm = p.grad.detach().data.norm(2)
total_grad_norm += param_grad_norm.item() ** 2
total_grad_norm = total_grad_norm ** 0.5
total_weight_norm = 0.0
for p in model.parameters():
param_weight_norm = p.detach().data.norm(2)
total_weight_norm += param_weight_norm.item() ** 2
total_weight_norm = total_weight_norm ** 0.5
print(total_grad_norm, total_weight_norm)
这个监测逻辑可以封装成函数,每隔一定步数打印一次。如果发现梯度范数在某个step之后呈指数增长,就不要继续训练,先停止并调整超参数,否则可能损坏模型权重。
二、学习率调整:从查找到调度
3D模型由于单批次样本数较少,梯度估计的方差比2D模型更大,因此学习率上限通常要设置得更保守。一个比较实用的方法是运行学习率查找,即从一个很小的学习率开始,每个mini-batch将学习率按比例增大,同时记录loss。当loss开始明显上升或出现发散迹象时,前一个学习率值可以作为后续训练的最大学习率参考。
学习率查找并非银弹,尤其对3D模型,loss曲线可能在一个较宽区间内都保持平坦,然后突然上升。这种情况下可以选择比拐点低一个数量级的值作为初始学习率,并配合预热机制逐步提升到目标值。预热对3D模型尤其重要,因为网络初始化阶段的梯度分布极不稳定,直接施加完整学习率很容易让前几层产生破坏性更新。
预热的实现很简单,一般在前500到2000个step内将学习率从0线性增加到目标学习率。预热结束后可以使用余弦退火逐步衰减学习率,让训练后期更新量不断减小,帮助模型收敛到更平坦的极小值。PyTorch提供了对应的调度器,组合使用时代码如下:
import torch
def warmup_lr_scheduler(optimizer, warmup_steps, target_lr):
def lr_lambda(step):
if step < warmup_steps:
return float(step) / float(max(1, warmup_steps))
return 1.0
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
warmup_scheduler = warmup_lr_scheduler(optimizer, warmup_steps=1000, target_lr=3e-4)
cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50000)
需要注意的是,代码中的比较运算符已经做了HTML转义处理,在页面中能正常显示。实际运行时Python源码里仍然是小于号。另外,如果使用自动混合精度,学习率的调整不受影响,但梯度裁剪需要放在scaler.unscale_之后。
逐层学习率是另一个可选策略。3D模型的主干网络通常在大规模2D数据集上预训练过,而3D分支是随机初始化。如果给所有层设置相同学习率,预训练层容易在初期被大梯度破坏。常见的做法是让3D分支使用较高学习率,例如3e-4,主干网络使用较低学习率,例如3e-5,比例控制在5到10倍之间。这样可以兼顾预训练知识的保留和新模块的快速学习。
三、优化器切换与超参数修正
Adam和AdamW是3D模型训练中最常用的优化器,但默认超参数并不总是适合3D任务。Adam的epsilon默认是1e-8,当梯度非常小或批次噪声较大时,二阶矩估计的数值可能不稳定,导致自适应学习率出现尖峰。将epsilon调大到1e-7或1e-6有时可以明显缓解发散,但会稍微降低收敛速度。AdamW解耦了权重衰减和梯度更新,对3D模型通常优于Adam加L2正则,因为L2正则与自适应学习率耦合时会削弱权重衰减的效果。
如果使用AdamW仍然发散,并且已经确认梯度范数没有异常,可以尝试切换到SGD加动量。SGD没有二阶矩自适应,更新幅度完全由学习率和梯度量级决定,虽然需要更仔细地调节学习率,但训练过程更可控,不容易出现自适应优化器那种突然的更新尖峰。对于3D模型,SGD搭配0.9动量、1e-4到3e-4的学习率,以及适度的权重衰减,往往能获得更稳定的收敛,只是训练周期可能需要更长。
另一个在大批次3D训练中值得考虑的是LAMB优化器。LAMB通过对每层参数进行归一化更新,能够在较大的批次下保持稳定,适合分布式3D模型训练场景。如果单卡显存限制导致被迫使用很小批次,LAMB的优势不明显,此时应优先保证梯度裁剪和预热。
梯度裁剪是防止发散的兜底手段。常用的裁剪方式包括按全局梯度范数裁剪和按参数梯度裁剪。对于3D模型,全局梯度范数裁剪到1.0是一个比较稳健的起点。如果训练任务对某些层的梯度量级特别敏感,可以改用按参数裁剪,但阈值设置不当可能引入额外偏差。PyTorch中两种方法分别对应torch.nn.utils.clip_grad_norm_和torch.nn.utils.clip_grad_value_,前者更常用。
优化器参数里还有一个容易被忽略的weight_decay。3D模型参数量大,过强的权重衰减会压制有效特征学习,过弱则容易过拟合。AdamW的weight_decay常见范围是0.01到0.1,如果训练发散,可以先检查是否把weight_decay设置得过大。过大的权重衰减会与高学习率叠加,产生类似发散的效果。
四、完整训练循环示例与监控
综合以上策略,一个稳定的3D模型训练循环应该包含学习率预热、余弦退火、梯度裁剪和必要的监控。下面给出一个基于PyTorch的简化示例,其中模型、数据加载器和损失函数都假设已经定义好:
import torch
model = build_3d_model()
criterion = torch.nn.MSELoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scaler = torch.cuda.amp.GradScaler()
warmup_steps = 1000
def lr_lambda(step):
if step < warmup_steps:
return float(step) / float(max(1, warmup_steps))
return 1.0
warmup_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50000)
global_step = 0
for epoch in range(num_epochs):
for batch in dataloader:
inputs, targets = batch
inputs = inputs.cuda()
targets = targets.cuda()
optimizer.zero_grad()
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
if global_step < warmup_steps:
warmup_scheduler.step()
else:
cosine_scheduler.step()
if global_step % 100 == 0:
print(loss.item())
global_step += 1
这段代码中,由于使用了自动混合精度,梯度裁剪必须在scaler.unscale_(optimizer)之后执行,否则裁剪的是缩放后的梯度,无法起到保护作用。学习率预热阶段使用线性增长,预热结束后切换为余弦退火,使总步数达到50000时学习率接近0。注意两个调度器不要同时执行,否则学习率会被重复更新。
在训练过程中除了打印损失值,还应定期保存模型检查点,尤其是在梯度范数、权重范数接近异常阈值时。可以设置一个自动停止条件,当损失连续若干步为NaN或超出合理数值范围时终止训练并回滚到最近的安全检查点。这样能避免发散后的无效计算,也方便后续复现问题。
如果发现训练在调整学习率后仍然发散,可以从数据端检查是否有异常标签或过大的体素强度。3D数据预处理中的归一化方式对训练稳定性也有影响,输入数值范围过大会导致第一层梯度异常。建议将输入归一化到均值为0、方差为1或0到1之间,必要时对体素强度做截断处理。
最后要强调的是,3D模型训练发散往往是多个因素共同作用的结果。单纯更换优化器可能掩盖问题,单纯降低学习率可能让训练陷入停滞。比较高效的做法是先从梯度范数和权重范数入手定位发散类型,再依次调整学习率调度、优化器超参数和梯度裁剪阈值。每一次只改动一个变量,记录损失曲线变化,这样才能建立对当前3D任务稳定训练区间的准确认识。