Adam梯度下降算法详解：从数学原理到PyTorch实战

1次阅读

没有评论

共计 2321 个字符，预计需要花费 6 分钟才能阅读完成。

梯度下降是深度学习模型训练的基石，它通过迭代调整参数来最小化损失函数。固定学习率会导致在平坦区域进展缓慢，而在陡峭区域容易震荡发散。自适应学习率算法正是为了解决这一矛盾而诞生的。

传统优化器各有明显缺陷：

SGD：学习率固定，难以适应不同参数的更新需求
Momentum：引入速度概念（$v_t = \beta v_{t-1} + (1-\beta)g_t$），但所有参数共享相同学习率
RMSprop：采用梯度平方的指数平均（$E[g^2]t = \gamma E[g^2] + (1-\gamma)g_t^2$），但缺少动量项

Adam 综合了两者优势：既保留动量项加速训练，又为每个参数维护自适应学习率。其更新规则可表示为：

$$
\theta_{t+1} = \theta_t – \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t
$$

Adam 同时计算梯度的一阶矩（均值）和二阶矩（未中心化的方差）：

$$
m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t \
v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2
$$

由于初始时刻的矩估计偏向零，需要进行修正：

$$
\hat{m}_t = \frac{m_t}{1-\beta_1^t} \
\hat{v}_t = \frac{v_t}{1-\beta_2^t}
$$

推荐初始设置：
– 学习率 η：3e-4（可上下浮动 10 倍）
– β1：0.9（控制一阶矩衰减）
– β2：0.999（控制二阶矩衰减）
– ε：1e-8（防止除零）

import torch
import torch.nn as nn
from torchvision import datasets, transforms

class CustomAdam:
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, max_grad_norm=1.0):
        self.params = list(params)
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.max_grad_norm = max_grad_norm
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]
        self.t = 0

    def step(self):
        self.t += 1
        for i, (param, m, v) in enumerate(zip(self.params, self.m, self.v)):
            if param.grad is None:
                continue

            grad = param.grad.data
            # 梯度裁剪
            if self.max_grad_norm > 0:
                grad = torch.clamp(grad, -self.max_grad_norm, self.max_grad_norm)

            # 更新一阶和二阶矩估计
            m.mul_(self.beta1).add_(grad, alpha=1-self.beta1)
            v.mul_(self.beta2).addcmul_(grad, grad, value=1-self.beta2)

            # 偏差修正
            m_hat = m / (1 - self.beta1**self.t)
            v_hat = v / (1 - self.beta2**self.t)

            # 参数更新
            param.data.addcdiv_(m_hat, v_hat.sqrt() + self.eps, value=-self.lr)

# MNIST 分类示例
model = nn.Sequential(nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)
opt = CustomAdam(model.parameters())

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.ToTensor()),
    batch_size=64, shuffle=True)

for epoch in range(5):
    for x, y in train_loader:
        x = x.view(-1, 784)
        out = model(x)
        loss = nn.CrossEntropyLoss()(out, y)

        opt.zero_grad()
        loss.backward()
        opt.step()