[TOC]

0. 基础

1. 导入

1
import torch.optim as optim

2. 常用的优化器

  • SGD/Momentum SGD
  • Adam/AdamW
  • AdaGrad
  • RMS prop

3. 使用框架

  • 生成优化器
1
2
3
4
import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)
  • 同一个模型定义不同的优化器
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
import torch.optim as optim

# params可以用字典传参，给不同的参数赋值不同的学习率
optim.SGD([{'params': model.base.parameters(), 'lr': 1e-2},
            {'params': model.classifier.parameters()}
        ], lr=1e-3, momentum=0.9)

# 给不同的参数设置不同的权重衰减系数，bias权重衰减设为0
bias_params = [p for name, p in model.named_parameters() if 'bias' in name]
others = [p for name, p in model.named_parameters() if 'bias' not in name]
optim.SGD([{'params': others},
           {'params': bias_params, 'weight_decay': 0}
           ], weight_decay=1e-2, lr=1e-2)
  • 使用优化器
1
2
3
4
5
6
7
8
9
10
for input, target in dataset:
    # 梯度清零
    optimizer.zero_grad()
    output = model(input)
    # 计算损失
    loss = loss_fn(output, target)
    # 损失反向传播（计算每个参数的梯度）
    loss.backward()
    # 根据梯度更新权重参数
    optimizer.step()

1. SGD/Momentum SGD

SGD：随机梯度优化，最简单常用的优化器，卷积神经网络时代比较好用，transformer时代被AdamW取代。

1.1 SGD原理

1. 计算梯度: ftf_t是模型函数，θt−1θ_{t-1}是上一步的参数，∇θ\nabla_θ是向量梯度计算，λ\lambda是权重衰减系数

gt=∇θft(θt−1)g_t=\nabla_θ f_t (θ_{t-1})

gt=gt+λθt−1g_t=g_t + \lambda θ_{t-1}

2. 梯度更新：λ\lambda是权重衰减系数，γ\gamma是学习率

θt=θt−1−γgt\theta_t = \theta_{t-1} -\gamma g_t

1.2 Momentum SGD原理

1. 计算梯度: 和SGD一样计算权重衰减后的梯度

gt=∇θft(θt−1)g_t=\nabla_θ f_t (θ_{t-1})

gt=gt+λθt−1g_t=g_t + \lambda θ_{t-1}

2. 计算动量：mtm_t是当前动量，mt−1m_{t-1}上一步的动量，β\beta是动量系数

m0=g0m_0 = g_0

mt=βmt−1+(1−β)gtm_t = \beta m_{t-1} + (1-\beta)g_t

目前新的版本使用1−τ1- \tau 代替1−β1- \beta，τ\tau作为动量抑制系数

mt=βmt−1+(1−τ)gtm_t = \beta m_{t-1} + (1-\tau)g_t

3. 梯度更新：γ\gamma是学习率

gt=btg_t = b_t

θt=θt−1−γgt\theta_t = \theta_{t-1} -\gamma g_t

1.3 代码

1
2
3
4
5
6
7
8
9
10
11
class SGD(Optimizer):
    def __init__(self, params, 
                 lr=1e-3, 
                 momentum=0, 
                 dampening=0,
                 weight_decay=0, 
                 nesterov=False, *, 
                 maximize: bool = False, 
                 foreach: Optional[bool] = None,
                 differentiable: bool = False, 
                 fused: Optional[bool] = None):

参数：

  • params：待优化的模型参数，通过**model.parameters()**获得
  • lr：学习率γ\gamma
  • momentum：梯度动量系数μ\mu
  • dampening：动量抑制τ\tau
  • eps：分母的添加项，增加数值稳定性1e-8
  • weight_decay：权重衰减系数，正则化系数λ\lambda

2. Adam/AdamW

2.1 Adam

1. 计算梯度: ftf_t是模型函数，θt−1θ_{t-1}是上一步的参数，∇θ\nabla_θ是向量梯度计算，λ\lambda是权重衰减系数

gt=∇θft(θt−1)g_t=\nabla_θ f_t (θ_{t-1})

gt=gt+λθt−1g_t=g_t + \lambda θ_{t-1}

2. 计算动量：mtm_t是当前动量，mt−1m_{t-1}上一步的动量，β1\beta_1是动量系数

mt=β1mt−1+(1−β1)gtm_t = \beta_1 m_{t-1} + (1-\beta_1)g_t

3. 计算二阶动量：btb_t是当前动量，bt−1b_{t-1}上一步的动量，μ\mu是动量系数

vt=β2vt−1+(1−β2)gt2v_t = \beta_2 v_{t-1} + (1-\beta_2)g^2_t

4. 快速启动：由于初始时$m_0=0, v_0=0,\beta \approx1 $，m和v更新很慢，需要很多步才能到正常值，需要进行偏差修正

mt^=mt/(1−β1t)\hat{m_t}=m_t/(1-\beta^t_1)

vt^=vt/(1−β2t)\hat{v_t}=v_t/(1-\beta^t_2)

5. 梯度更新：γ\gamma是学习率

θt=θt−1−γmt^/(v^t+ϵ)\theta_t = \theta_{t-1} -\gamma \hat{m_t}/(\sqrt {\hat v_t} + \epsilon)

2.2 AdamW

AdamW是对Adam的修正，修改了权重衰减的位置，权重衰减放到了最后的损失位置，而不是最开始的梯度位置。

1. 计算梯度: ftf_t是模型函数，θt−1θ_{t-1}是上一步的参数，∇θ\nabla_θ是向量梯度计算，λ\lambda是权重衰减系数

gt=∇θft(θt−1)g_t=\nabla_θ f_t (θ_{t-1})

2. 计算动量：mtm_t是当前动量，mt−1m_{t-1}上一步的动量，β1\beta_1是动量系数

mt=β1mt−1+(1−β1)gtm_t = \beta_1 m_{t-1} + (1-\beta_1)g_t

3. 计算二阶动量：btb_t是当前动量，bt−1b_{t-1}上一步的动量，μ\mu是动量系数

vt=β2vt−1+(1−β2)gt2v_t = \beta_2 v_{t-1} + (1-\beta_2)g^2_t

4. 快速启动：由于初始时$m_0=0, v_0=0,\beta \approx1 $，m和v更新很慢，需要很多步才能到正常值，需要进行偏差修正

mt^=mt/(1−β1t)\hat{m_t}=m_t/(1-\beta^t_1)

vt^=vt/(1−β2t)\hat{v_t}=v_t/(1-\beta^t_2)

5. 梯度更新：γ\gamma是学习率

θt=θt−1−γλθt−1\theta_t=\theta_{t-1} - \gamma\lambda θ_{t-1}

θt=θt−1−γmt^/(v^t+ϵ)\theta_t = \theta_{t-1} -\gamma \hat{m_t}/(\sqrt {\hat v_t} + \epsilon)

3. AdaGrad/RMS prop

相当于动量系数β1\beta_1为0的Adam，根据参数之前的梯度作为当前参数的权重

4. Muon