问题描述
- 对 stable diffusion 3 进行 ip-adapter 微调,正常训练 2 w 步后,loss 出现不稳定状态,并出现 Not a number
- 问题定位:由于 loss 并没有变成无限大(梯度爆炸),那么应该是梯度消失。
解决方案
- 降低学习率
- 参考 huggingface 官方论坛1,其中有用户提到:“我也遇到过几次了。就我而言,我能够通过降低学习率来解决这个问题,但你的学习率已经很低了,所以不幸的是,我不太确定。”
- 该博客中的学习率已经是 learning_rate=1e-6,而本文的情况是 1e-4 会在 200步 nan,8e-05 会在 600步 nan,降到 1e-5 稳定训练很长时间后,一觉醒来,发现在 2 w 步之后出现 nan 🫠
- 自动调节学习率
自动调节学习率的方式很多2,本文先从最简单地尝试起来,即“线性调节”
(1)先 warm up,在前 1/10 个训练步中,学习率从 0 调节到设定的 1e-5
(2)再衰减,在后续的训练步中,线性地从 1e-5 逐渐减为 0
线性调节学习率代码
# S-TODO 学习率调节器 lr_scheduler refer to https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.LambdaLR.html#torch.optim.lr_scheduler.LambdaLR
from torch.optim.lr_scheduler import LambdaLR
def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
def lr_lambda(current_step):
if current_step < num_warmup_steps:
return float(current_step) / float(max(1.0, num_warmup_steps))
return max(
0.0, float(num_training_steps - current_step) / float(max(1.0, num_training_steps - num_warmup_steps))
)
return LambdaLR(optimizer, lr_lambda, last_epoch)
...
optimizer = torch.optim.AdamW(xxx)
# dataloader
train_dataset = MyDataset(xxx)
train_dataloader = torch.utils.data.DataLoader(xxx)
# S-TODO Define the 学习率 的 scheduler
num_training_steps = len(train_dataloader) * args.num_train_epochs
num_warmup_steps = int(0.1 * num_training_steps) # 10% warmup
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)
...
训练代码
...
# Backpropagate
accelerator.backward(loss)
optimizer.step()
# 更新下一步的学习率
scheduler.step()
optimizer.zero_grad()
...