pytorch优化器动态学习率和动量设置（scheduler & momentum）

一、动量（momentum）

　可以给优化器加上一个动量，可以有效缓解局部最优问题。

　原理上就是让优化过程从

　　W = W - lr * dW

　变成

　　V = momentum * V - lr * dW

　　W = W + V

　使用示例：

from torch import optim
...
model = Model()
optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.9)
...

 

二、优化器动态学习率设置（scheduler）

　可以让学习率随着epoch的增大而减小，此处以ExponentialLR为例

　使用示例：

from torch.optim.lr_scheduler import ExponentialLR
...
optimizer = optim.SGD(catp.parameters(), lr=0.005, momentum=0.9)
scheduler = ExponentialLR(optimizer, gamma=0.99)
for epoch in range(epochs):
    for i, batch_data in enumerate(dataloader):
        loss = ...
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    scheduler.step()
...

 

三、带有幂(exp)衰减因子的退火余弦热重启(CosineAnnealingWarmRestarts)的scheduler

def exp_cos_annealing_warm_restart(epoch_index: int, optimizer: Optimizer, cos_scheduler: CosineAnnealingWarmRestarts, gamma: float):
    cos_scheduler.step()
    optimizer.param_groups[-1]['lr'] = gamma ** epoch_index * optimizer.param_groups[-1]['lr']


# training ...

optimizer = optim.Adam(params=model.parameters(), lr=1e-4, betas=(0.9, 0.999), weight_decay=0.01,
                               amsgrad=True)
scheduler = CosineAnnealingWarmRestarts(optimizer=optimizer, T_0=2, T_mult=2, eta_min=5e-6)
...
for epoch_idx in range(100):
    model.train()
    for batch_idx, (data, label) in enumerate(train_loader):
        optimizer.zero_grad()
        loss = loss_fn(model(data), label)
        loss.backward()
        optimizer.step()
        exp_cos_annealing_warm_restart(epoch_index=epoch_idx + 1, optimizer=optimizer, cos_scheduler=scheduler, gamma=0.99)