优化器、学习率与正则化

"""Epoch / Batch size / Iteration 三个数到底谁决定谁。

这三个词天天出现在训练日志里，但很多人说不清它们的换算关系。
本文件不依赖 torch，纯标准库，可以直接跑出确定的数字自己核对。

换算关系只有一条：
    每个 epoch 的 iteration 数 = ceil(样本总数 / batch_size)
「ceil」是向上取整——最后一个不满的批次也要算一次参数更新。
"""
import math


def plan(sample_num, batch_size, epochs):
    """算出一次训练总共要做多少次参数更新。"""
    iter_per_epoch = math.ceil(sample_num / batch_size)
    last_batch = sample_num - (iter_per_epoch - 1) * batch_size
    return {
        'sample_num': sample_num,
        'batch_size': batch_size,
        'epochs': epochs,
        'iter_per_epoch': iter_per_epoch,
        'last_batch_size': last_batch,        # 最后一批可能不满
        'total_iter': iter_per_epoch * epochs,
    }


def show(info):
    print('样本 %d 条，batch_size=%d，共 %d 轮'
          % (info['sample_num'], info['batch_size'], info['epochs']))
    print('  每轮 iteration 数 = ceil(%d / %d) = %d'
          % (info['sample_num'], info['batch_size'], info['iter_per_epoch']))
    print('  最后一批只有 %d 条样本' % info['last_batch_size'])
    print('  全程参数更新次数 = %d × %d = %d'
          % (info['iter_per_epoch'], info['epochs'], info['total_iter']))
    print()


# ---------------------------------------------------------------- 三种梯度下降
# 区别只在 batch_size 取多少，公式本身一个字没变
VARIANTS = [
    # 名称,           batch_size 取法,        一轮更新几次,        特点
    ('全批量 BGD', '等于全部样本数', '1 次', '方向最准，但一步要算完所有样本，慢且吃内存'),
    ('随机 SGD', '1', '样本数那么多次', '更新快、噪声大，损失曲线剧烈抖动'),
    ('小批量 Mini-Batch', '8 / 16 / 32 / 64…', 'ceil(N/B) 次', '实际训练唯一在用的方案，噪声与效率的折中'),
]


def compare_variants(sample_num=50000):
    print('样本总数 %d 时，三种梯度下降一轮各更新几次：' % sample_num)
    print('  全批量  : 1 次')
    print('  随机    : %d 次' % sample_num)
    for b in (8, 256):
        print('  小批量 B=%-4d: %d 次' % (b, math.ceil(sample_num / b)))
    print()


if __name__ == '__main__':
    compare_variants()

    # 一个能手算核对的例子：50000 张图，batch_size=256
    # 50000 / 256 = 195.3125 -> 向上取整 196
    show(plan(sample_num=50000, batch_size=256, epochs=10))

    # 不整除的小例子，方便一眼看清「最后一批不满」
    show(plan(sample_num=2000, batch_size=8, epochs=50))

    for name, how, times, note in VARIANTS:
        print('%-18s batch_size %-16s 一轮 %-14s %s' % (name, how, times, note))

"""过拟合 / 欠拟合的判别器：只看训练与验证两条损失曲线的关系。

不依赖 torch，把训练过程中记录下来的两条曲线喂进来就能出诊断结论。
真实项目里把 train_loss / valid_loss 每轮 append 进列表，训练完直接调用。

判别规则（三种典型形态）：
    1. 两条都高、都还在降        -> 欠拟合：模型太小或训练轮数不够
    2. 训练一直降、验证先降后升  -> 过拟合：该上正则化了
    3. 两条都低且贴在一起        -> 拟合良好
"""


def diagnose(train_loss, valid_loss, gap_ratio=0.3, patience=5):
    """返回 (结论, 解释)。

    gap_ratio: 验证损失比训练损失高出多少比例算「差距过大」
    patience : 验证损失连续多少轮不再下降算「已经转折」
    """
    assert len(train_loss) == len(valid_loss), '两条曲线长度必须一致'
    assert len(train_loss) >= patience + 1, '轮数太少，看不出趋势'

    last_train = train_loss[-1]
    last_valid = valid_loss[-1]
    best_valid = min(valid_loss)
    best_epoch = valid_loss.index(best_valid)

    # 训练损失还在明显下降吗（拿最后 patience 轮和之前对比）
    train_still_falling = train_loss[-patience - 1] - last_train > 1e-4

    # 验证损失从最低点之后又涨回去了多少
    valid_rebound = last_valid - best_valid

    if last_train > 0.6 * train_loss[0] and train_still_falling:
        return ('欠拟合',
                '训练损失本身还很高且仍在下降，说明模型连训练集都没学明白：'
                '加深网络、加宽隐藏层、延长训练轮数、或者把学习率调大一点。')

    if valid_rebound > 1e-4 and best_epoch < len(valid_loss) - patience:
        return ('过拟合',
                '验证损失在第 %d 轮触底 %.4f，之后反弹到 %.4f，而训练损失还在降——'
                '模型开始背答案了：上 Dropout / weight_decay，或直接早停到第 %d 轮。'
                % (best_epoch + 1, best_valid, last_valid, best_epoch + 1))

    if last_valid - last_train > gap_ratio * max(last_train, 1e-8):
        return ('轻度过拟合',
                '验证损失 %.4f 明显高于训练损失 %.4f，差距已经拉开，'
                '虽然还没反弹，但正则化可以先加上了。' % (last_valid, last_train))

    return ('拟合良好',
            '两条曲线都降到低位且贴得很近，训练 %.4f / 验证 %.4f。' % (last_train, last_valid))


# ---------------------------------------------------------------- 三组造出来的曲线
# 这些数字是手工编排的典型形态，用来验证判别逻辑本身，不是某次真实训练的结果
UNDERFIT_TRAIN = [1.38, 1.35, 1.33, 1.31, 1.29, 1.27, 1.25, 1.23]
UNDERFIT_VALID = [1.39, 1.36, 1.34, 1.32, 1.30, 1.28, 1.26, 1.24]

OVERFIT_TRAIN = [1.40, 0.95, 0.62, 0.40, 0.25, 0.15, 0.09, 0.05, 0.03, 0.02]
OVERFIT_VALID = [1.42, 1.00, 0.72, 0.61, 0.58, 0.60, 0.66, 0.74, 0.83, 0.91]

GOOD_TRAIN = [1.40, 0.90, 0.60, 0.45, 0.38, 0.34, 0.32, 0.31, 0.30, 0.30]
GOOD_VALID = [1.43, 0.94, 0.64, 0.48, 0.41, 0.36, 0.34, 0.33, 0.32, 0.32]


def ascii_curve(train, valid, width=48):
    """不装 matplotlib 也能看走势：把两条曲线画成字符图。"""
    hi = max(max(train), max(valid))
    lo = min(min(train), min(valid))
    span = (hi - lo) or 1.0
    for i, (t, v) in enumerate(zip(train, valid)):
        tp = int((t - lo) / span * (width - 1))
        vp = int((v - lo) / span * (width - 1))
        row = [' '] * width
        row[tp] = 'T'                       # T = 训练损失
        row[vp] = 'V' if vp != tp else '*'  # V = 验证损失，重合画 *
        print('  epoch %2d |%s|  train %.3f  valid %.3f'
              % (i + 1, ''.join(row), t, v))


if __name__ == '__main__':
    for name, tr, va in (('一、欠拟合', UNDERFIT_TRAIN, UNDERFIT_VALID),
                         ('二、过拟合', OVERFIT_TRAIN, OVERFIT_VALID),
                         ('三、拟合良好', GOOD_TRAIN, GOOD_VALID)):
        print(name)
        ascii_curve(tr, va)
        verdict, why = diagnose(tr, va)
        print('  诊断结论：%s —— %s\n' % (verdict, why))

"""指数加权平均：Momentum、RMSProp、Adam 的共同零件。

一句话：新值 = 上一次的平均 × β + 当前值 × (1-β)。
β 越大，曲线越平滑，越不理会当前这一下的抖动。
"""
import torch
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

N = 30


def ewa(beta):
    torch.manual_seed(0)
    temperature = torch.randn(size=[N]) * 10     # 造 30 天的随机气温
    avg = []
    for idx, temp in enumerate(temperature, 1):
        if idx == 1:
            avg.append(temp)                      # 第一天没有历史，等于自身
            continue
        avg.append(avg[idx - 2] * beta + (1 - beta) * temp)
    return temperature, avg


if __name__ == '__main__':
    days = torch.arange(1, N + 1)
    for beta in (0.5, 0.9):
        temperature, avg = ewa(beta)
        plt.scatter(days, temperature)            # 散点是原始值
        plt.plot(days, avg, color='r')            # 折线是加权平均
        plt.title('β = %.1f' % beta)
        plt.show()

    # 把同一批数据的两条曲线打印出来对比末值
    print('β=0.5 末值 %.4f' % ewa(0.5)[1][-1])
    print('β=0.9 末值 %.4f  ← 更平缓，更依赖历史' % ewa(0.9)[1][-1])

"""把四种优化器的更新公式逐字翻译成 Python，不依赖任何第三方库。

为什么要手写一遍：`torch.optim.Adam(...)` 这一行把公式全藏起来了，
手写版每一步的中间量都打印出来，能和纸上推导逐位对账。

统一实验条件（四种方法完全可比）：
    损失  loss = w² / 2      ->  梯度 g = w
    起点  w = 1.0            ->  第一次的梯度必定是 1.0
    lr    = 0.01
"""
import math

LR = 0.01


def sgd(w, steps):
    """最朴素的梯度下降：w <- w - lr * g，没有任何记忆。"""
    for i in range(1, steps + 1):
        g = w                                  # loss = w²/2 的导数就是 w
        w = w - LR * g
        yield i, g, w, {}


def momentum(w, steps, beta=0.9):
    """动量法。注意这里用的是 PyTorch 的实现口径：

        buf <- beta * buf + g          （不乘 (1-beta)）
        w   <- w - lr * buf

    教材里常写成指数加权平均 s <- beta*s + (1-beta)*g，
    两者只差一个常数倍，把它并进学习率即可，收敛行为一致。
    """
    buf = 0.0
    for i in range(1, steps + 1):
        g = w
        buf = beta * buf + g                   # 历史梯度攒下来的「惯性」
        w = w - LR * buf
        yield i, g, w, {'buf': buf}


def adagrad(w, steps, eps=1e-10):
    """AdaGrad：把历史梯度的平方**求和**，用它给学习率做分母。

    分母只增不减 -> 学习率单调下降 -> 后期可能过早走不动。
    """
    s = 0.0
    for i in range(1, steps + 1):
        g = w
        s = s + g * g                          # 平方和，永远在变大
        w = w - LR * g / (math.sqrt(s) + eps)
        yield i, g, w, {'s': s, 'lr_eff': LR / (math.sqrt(s) + eps)}


def rmsprop(w, steps, alpha=0.9, eps=1e-8):
    """RMSProp：把 AdaGrad 的「求和」换成「指数加权平均」。

    分母不再单调变大，久远的梯度会被逐渐遗忘，解决过早熄火。
    PyTorch 里这个系数叫 alpha，就是公式里的 beta。
    """
    s = 0.0
    for i in range(1, steps + 1):
        g = w
        s = alpha * s + (1 - alpha) * g * g    # 指数加权平均，会遗忘
        w = w - LR * g / (math.sqrt(s) + eps)
        yield i, g, w, {'s': s, 'lr_eff': LR / (math.sqrt(s) + eps)}


def adam(w, steps, beta1=0.9, beta2=0.99, eps=1e-8):
    """Adam = Momentum（一阶矩 m）+ RMSProp（二阶矩 v）+ 偏差修正。

    偏差修正那两下除法是必需的：m 和 v 都从 0 起步，
    前几步会被严重低估，除以 (1 - beta^t) 正好把它拉回真实量级。
    """
    m = v = 0.0
    for i in range(1, steps + 1):
        g = w
        m = beta1 * m + (1 - beta1) * g        # 一阶矩：梯度的方向
        v = beta2 * v + (1 - beta2) * g * g    # 二阶矩：梯度的幅度
        m_hat = m / (1 - beta1 ** i)           # 偏差修正
        v_hat = v / (1 - beta2 ** i)
        w = w - LR * m_hat / (math.sqrt(v_hat) + eps)
        yield i, g, w, {'m': m, 'v': v, 'm_hat': m_hat, 'v_hat': v_hat}


def run(name, gen):
    print('== %s ==' % name)
    for i, g, w, extra in gen:
        detail = '  '.join('%s=%.6f' % kv for kv in extra.items())
        print('  第%d次  梯度 %.6f  更新后 w %.6f   %s' % (i, g, w, detail))
    print()


if __name__ == '__main__':
    run('SGD', sgd(1.0, 2))
    run('Momentum', momentum(1.0, 2))
    run('AdaGrad', adagrad(1.0, 2))
    run('RMSProp', rmsprop(1.0, 2))
    run('Adam', adam(1.0, 2))

    # 跑长一点，看四种方法 30 步之后各自走到哪里
    print('== 30 步之后的 w（越接近 0 越好）==')
    for name, fn in (('SGD', sgd), ('Momentum', momentum),
                     ('AdaGrad', adagrad), ('RMSProp', rmsprop), ('Adam', adam)):
        last = 1.0
        for _, _, last, _ in fn(1.0, 30):
            pass
        print('  %-10s w = %.6f' % (name, last))

"""四种优化器在同一个损失函数上各走两步，看更新幅度的差别。

损失统一取 loss = w² / 2，它对 w 的导数就是 w 本身，
所以起点 w=1.0 时第一次的梯度必定是 1.0，四种方法完全可比。
"""
import torch


def run(name, make_optimizer):
    w = torch.tensor([1.0], requires_grad=True)
    optimizer = make_optimizer(w)

    for step in (1, 2):
        loss = ((w ** 2) / 2.0).sum()      # 每步都要重新前向，w 变了
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print('%-10s 第%d次  梯度 %.6f  更新后 w %.6f'
              % (name, step, w.grad.item(), w.item()))
    print()


if __name__ == '__main__':
    # SGD + momentum：带惯性，第二步的位移明显比纯 SGD 大
    run('Momentum', lambda w: torch.optim.SGD([w], lr=0.01, momentum=0.9))

    # AdaGrad：累积平方梯度，学习率只降不升，后期容易走不动
    run('AdaGrad', lambda w: torch.optim.Adagrad([w], lr=0.01))

    # RMSProp：把 AdaGrad 的「求和」换成「指数加权平均」，解决过早熄火
    #          alpha 就是公式里的 β
    run('RMSProp', lambda w: torch.optim.RMSprop([w], lr=0.01, alpha=0.9))

    # Adam：Momentum（一阶矩）+ RMSProp（二阶矩），默认首选
    run('Adam', lambda w: torch.optim.Adam([w], lr=0.01, betas=(0.9, 0.99)))

"""三种学习率衰减策略：调度器每个 epoch 走一步，不是每个 batch。"""
import torch
from torch import optim
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

LR, ITERATION, MAX_EPOCH = 0.1, 10, 200


def collect(make_scheduler, title):
    w = torch.tensor([1.0], requires_grad=True)
    x = torch.tensor([1.0])
    y_true = torch.tensor([0.0])

    optimizer = optim.SGD([w], lr=LR, momentum=0.9)
    scheduler = make_scheduler(optimizer)

    lr_list, epoch_list = [], []
    for epoch in range(MAX_EPOCH):
        lr_list.append(scheduler.get_last_lr()[0])
        epoch_list.append(epoch)

        for _ in range(ITERATION):                 # 一个 epoch 里的每个 batch
            loss = ((w * x - y_true) ** 2) / 2.0
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        scheduler.step()      # 关键：放在 epoch 循环里，不要放进 batch 循环

    plt.plot(epoch_list, lr_list, label=title)
    plt.xlabel('Epoch')
    plt.ylabel('Learning rate')
    plt.legend()
    plt.show()
    print('%s：起始 %.4f -> 结束 %.6f' % (title, lr_list[0], lr_list[-1]))


if __name__ == '__main__':
    # 等间隔：每 50 个 epoch 砍一半
    collect(lambda o: optim.lr_scheduler.StepLR(o, step_size=50, gamma=0.5),
            '等间隔衰减 StepLR')

    # 指定间隔：只在这几个点上砍
    collect(lambda o: optim.lr_scheduler.MultiStepLR(o, milestones=[50, 125, 160],
                                                     gamma=0.5),
            '指定间隔衰减 MultiStepLR')

    # 指数：每个 epoch 都乘 gamma，平滑下降
    collect(lambda o: optim.lr_scheduler.ExponentialLR(o, gamma=0.95),
            '指数衰减 ExponentialLR')

"""Dropout 与 BatchNorm 的行为：注意训练模式和测试模式不一样。"""
import torch
import torch.nn as nn

torch.manual_seed(0)

# ---------------- Dropout ----------------
dropout = nn.Dropout(p=0.4)
inputs = torch.randint(0, 10, size=[1, 4]).float()
layer = nn.Linear(4, 5)
y = torch.relu(layer(inputs))
print('未失活：', y)

dropout.train()                 # 训练模式：随机置 0，剩下的按 1/(1-p) 放大
print('训练时失活：', dropout(y))

dropout.eval()                  # 测试模式：Dropout 直接变成恒等映射
print('测试时不失活：', dropout(y))

# ---------------- BatchNorm ----------------
# 参数 2 = 通道数；affine=True 表示缩放 γ 和平移 β 是可学习的
bn = nn.BatchNorm2d(2, eps=1e-05, momentum=0.1, affine=True)
x = torch.randn(1, 2, 3, 4)     # (batch, channel, height, width)
out = bn(x)
print('BN 输出形状', out.shape)  # 形状不变，只是把每个通道的分布拉正
print('可学习的缩放 γ', bn.weight.data)   # 初始全 1
print('可学习的平移 β', bn.bias.data)     # 初始全 0

"""model.train() 与 model.eval() 到底切换了什么。

这两个方法不改任何权重、不算任何梯度，它们只做一件事：
把网络里所有子模块的 self.training 标志位设成 True / False。
真正因此改变行为的只有两类层：

    Dropout    训练时随机置 0 并按 1/(1-p) 放大；测试时完全不动
    BatchNorm  训练时用**当前这批**的均值方差，同时更新全局running统计量；
               测试时改用累计下来的 running_mean / running_var

忘了切换的后果：
    忘了 eval()  -> 预测结果每次都不一样，且准确率虚低
    忘了 train() -> Dropout 永远不生效，正则化等于没加
"""
import torch
import torch.nn as nn

torch.manual_seed(0)


def show_dropout_behavior():
    print('===== Dropout =====')
    dropout = nn.Dropout(p=0.5)
    x = torch.ones(1, 8)                     # 全 1，方便看出哪些被置 0、哪些被放大

    dropout.train()
    print('训练模式连打三次（每次都不同）：')
    for _ in range(3):
        print('  ', dropout(x))
    # 存活的元素值是 1/(1-0.5) = 2.0，不是 1.0
    # 放大的目的：让这一层输出的**期望总量**和不失活时一致，
    # 否则测试时不失活，数值量级会突然翻倍。

    dropout.eval()
    print('测试模式连打三次（完全相同，就是恒等映射）：')
    for _ in range(3):
        print('  ', dropout(x))


def show_bn_behavior():
    print('\n===== BatchNorm =====')
    bn = nn.BatchNorm1d(4)
    x = torch.tensor([[1.0, 2.0, 3.0, 4.0],
                      [5.0, 6.0, 7.0, 8.0]])

    print('初始 running_mean =', bn.running_mean)   # 全 0
    print('初始 running_var  =', bn.running_var)    # 全 1

    bn.train()
    out_train = bn(x)
    print('训练模式输出（用这一批自己的均值方差标准化）：\n', out_train)
    # 每列两个数标准化后必然是 -1 和 +1（两样本情形），与列的绝对大小无关
    print('更新后的 running_mean =', bn.running_mean)
    # running_mean = (1 - momentum) * 旧值 + momentum * 本批均值
    #              = 0.9 * 0 + 0.1 * 本批均值

    bn.eval()
    out_eval = bn(x)
    print('测试模式输出（改用 running 统计量，数值完全不同）：\n', out_eval)

    print('可学习的缩放 γ =', bn.weight.data)        # 初始全 1
    print('可学习的平移 β =', bn.bias.data)          # 初始全 0


def show_batch_size_one_trap():
    print('\n===== BN 的 batch_size=1 陷阱 =====')
    bn = nn.BatchNorm1d(4)
    bn.train()
    one_sample = torch.randn(1, 4)
    try:
        bn(one_sample)
    except Exception as exc:
        print('训练模式下喂 1 条样本：', type(exc).__name__, exc)
    # 原因：一条样本算不出方差（分母为 0），BatchNorm1d 直接拒绝。
    # 现实中触发场景：最后一个 batch 恰好只剩 1 条 ->
    # DataLoader 加 drop_last=True 即可绕开。

    bn.eval()
    print('测试模式下喂 1 条样本没问题：', bn(one_sample).shape)


class Net(nn.Module):
    """一个同时含 Dropout 和 BN 的网络，用来演示整网切换。"""

    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 16)
        self.bn = nn.BatchNorm1d(16)
        self.drop = nn.Dropout(p=0.3)
        self.fc2 = nn.Linear(16, 2)

    def forward(self, x):
        x = torch.relu(self.bn(self.fc1(x)))
        x = self.drop(x)
        return self.fc2(x)


def show_whole_net_switch():
    print('\n===== 整网一次性切换 =====')
    net = Net()
    x = torch.randn(4, 8)

    net.train()
    a, b = net(x), net(x)
    print('train() 下两次前向是否相同：', torch.equal(a, b))   # False

    net.eval()
    with torch.no_grad():
        c, d = net(x), net(x)
    print('eval() 下两次前向是否相同：', torch.equal(c, d))    # True
    # 结论：预测阶段必须 eval()，否则同一张图连问两次会给出两个答案。


if __name__ == '__main__':
    show_dropout_behavior()
    show_bn_behavior()
    show_batch_size_one_trap()
    show_whole_net_switch()

"""L1 / L2 惩罚项：给损失函数加一条「别把权重养太大」的约束。

两种写法：
  1. 手写进损失里     —— 看得见公式，L1 只能这么写
  2. 交给 weight_decay —— 一个参数搞定，等价于 L2，实际项目里用这个

注意 weight_decay 同时作用在 weight 和 bias 上；严格做法是把 bias
和 BatchNorm 的参数排除在外，见文件末尾的分组写法。
"""
import torch
from torch import nn, optim


def l1_penalty(model):
    """L1：所有参数绝对值之和。倾向于把不重要的权重压成正好 0（稀疏）。"""
    return sum(p.abs().sum() for p in model.parameters())


def l2_penalty(model):
    """L2：所有参数平方和。倾向于把权重整体压小，但不会压到 0。"""
    return sum((p ** 2).sum() for p in model.parameters())


def train_with_manual_penalty(model, dataloader, kind='l2',
                              lam=1e-4, lr=1e-3, epochs=5):
    """手写惩罚项版本：损失 = 原始损失 + lam × 惩罚。"""
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)   # 注意不填 weight_decay

    for epoch in range(epochs):
        model.train()
        for bx, by in dataloader:
            base_loss = criterion(model(bx), by)
            penalty = l1_penalty(model) if kind == 'l1' else l2_penalty(model)
            loss = base_loss + lam * penalty        # lam 就是惩罚的力度

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print('epoch %d  原始损失 %.5f  惩罚项 %.5f'
              % (epoch + 1, base_loss.item(), penalty.item()))
    return model


def build_optimizer_with_weight_decay(model, lr=1e-3, wd=1e-4):
    """框架版：weight_decay 就是 L2，写一个参数即可，不用改损失函数。"""
    return optim.Adam(model.parameters(), lr=lr, weight_decay=wd)


def build_optimizer_no_decay_on_bias(model, lr=1e-3, wd=1e-4):
    """更严谨的分组：bias 和 BatchNorm 的参数不做衰减。

    这两类参数本来就该自由取值，压小它们只会拖慢收敛，
    对准确率没有好处。参数量大的模型上这个差别很明显。
    """
    decay, no_decay = [], []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # bias 与 BN 的 weight/bias 都是一维的，据此区分最省事
        if param.ndim <= 1 or name.endswith('.bias'):
            no_decay.append(param)
        else:
            decay.append(param)

    print('参与衰减的参数张量 %d 个，豁免的 %d 个' % (len(decay), len(no_decay)))
    return optim.Adam([
        {'params': decay, 'weight_decay': wd},
        {'params': no_decay, 'weight_decay': 0.0},
    ], lr=lr)


if __name__ == '__main__':
    net = nn.Sequential(
        nn.Linear(20, 128), nn.ReLU(),
        nn.BatchNorm1d(128),
        nn.Linear(128, 4),
    )

    # 两种惩罚在同一组参数上的数值差别：L2 对大权重罚得更狠（平方）
    print('L1 惩罚 =', l1_penalty(net).item())
    print('L2 惩罚 =', l2_penalty(net).item())

    build_optimizer_with_weight_decay(net)
    build_optimizer_no_decay_on_bias(net)

"""早停（Early Stopping）：验证损失不再下降就收手，并回滚到最好的那一版。

早停是最便宜的正则化手段——不改网络、不改损失函数、不加超参数，
只是「别再练了」。但它有两个新手常漏的细节：

  1. 判据必须是**验证集**指标，拿训练损失判早停等于没判；
  2. 停下来之后要把权重**回滚到最佳轮次**，而不是留着最后一轮的权重——
     最后一轮已经是开始变差的那一版了。

这个类不依赖 torch，可以直接跑；save/restore 两个钩子留给调用方注入。
"""


class EarlyStopping:
    def __init__(self, patience=5, min_delta=1e-4, mode='min'):
        """patience : 连续多少轮没有改善就停
        min_delta: 改善多少才算数，防止在噪声里反复横跳
        mode     : 'min' 盯损失越小越好，'max' 盯准确率越大越好
        """
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.best = None
        self.best_epoch = -1
        self.bad_epochs = 0
        self.should_stop = False

    def _is_better(self, value):
        if self.best is None:
            return True
        if self.mode == 'min':
            return value < self.best - self.min_delta
        return value > self.best + self.min_delta

    def step(self, value, epoch, on_improve=None):
        """每个 epoch 算完验证指标后调用一次。返回是否应当停止。"""
        if self._is_better(value):
            self.best = value
            self.best_epoch = epoch
            self.bad_epochs = 0
            if on_improve is not None:
                on_improve()               # 这里把当前权重存下来
        else:
            self.bad_epochs += 1
            if self.bad_epochs >= self.patience:
                self.should_stop = True
        return self.should_stop


def train_with_early_stopping(model, train_loader, valid_loader,
                              criterion, optimizer, max_epochs=200,
                              patience=5, ckpt='model/best.pt'):
    """接进真实训练循环的样子。"""
    import copy
    import torch

    stopper = EarlyStopping(patience=patience, mode='min')
    best_state = None

    for epoch in range(max_epochs):
        # ---- 训练一轮
        model.train()                       # 打开 Dropout / BN 的训练行为
        for bx, by in train_loader:
            loss = criterion(model(bx), by)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # ---- 在验证集上评估
        model.eval()                        # 关掉 Dropout，BN 改用全局统计量
        total, num = 0.0, 0
        with torch.no_grad():               # 评估不需要梯度
            for bx, by in valid_loader:
                total += criterion(model(bx), by).item() * len(by)
                num += len(by)
        valid_loss = total / num
        print('epoch %3d  valid_loss %.5f' % (epoch + 1, valid_loss))

        def keep():
            nonlocal best_state
            best_state = copy.deepcopy(model.state_dict())

        if stopper.step(valid_loss, epoch, on_improve=keep):
            print('验证损失连续 %d 轮没有改善，停在第 %d 轮' % (patience, epoch + 1))
            break

    # ---- 回滚到最好的那一版，这一步不能省
    if best_state is not None:
        model.load_state_dict(best_state)
        torch.save(best_state, ckpt)
        print('已回滚到第 %d 轮的权重，验证损失 %.5f'
              % (stopper.best_epoch + 1, stopper.best))
    return model


if __name__ == '__main__':
    # 用 overfit_diagnose.py 里那条过拟合曲线走一遍，看早停停在哪
    valid_curve = [1.42, 1.00, 0.72, 0.61, 0.58, 0.60, 0.66, 0.74, 0.83, 0.91]
    stopper = EarlyStopping(patience=3, mode='min')
    for epoch, v in enumerate(valid_curve):
        stopped = stopper.step(v, epoch)
        print('epoch %2d  valid %.2f  best %.2f(第%d轮)  连续变差 %d 轮  %s'
              % (epoch + 1, v, stopper.best, stopper.best_epoch + 1,
                 stopper.bad_epochs, '-> 停止' if stopped else ''))
        if stopped:
            break

"""训练循环上的三个插槽：换优化器、接调度器、插正则化层。

五步主循环（前向 -> 算损失 -> 清零 -> 反向 -> 更新）一个字都不用改，
调优要动的只有这三处。把这份文件当成「改哪里」的对照表。
"""
import torch
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset


# ---------------------------------------------------------------- 插槽一：优化器
def build_optimizer(model, kind='adam', lr=1e-3, weight_decay=0.0):
    """换优化器只改这一行，训练循环本身不动。"""
    if kind == 'sgd':
        return optim.SGD(model.parameters(), lr=lr,
                         weight_decay=weight_decay)
    if kind == 'momentum':
        # 加惯性：几乎没有理由用不带动量的纯 SGD
        return optim.SGD(model.parameters(), lr=lr, momentum=0.9,
                         weight_decay=weight_decay)
    if kind == 'rmsprop':
        return optim.RMSprop(model.parameters(), lr=lr, alpha=0.9,
                             weight_decay=weight_decay)
    # 默认首选；weight_decay 就是 L2 正则，填上即可，不用改损失函数
    return optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999),
                      weight_decay=weight_decay)


# ---------------------------------------------------------------- 插槽二：调度器
def build_scheduler(optimizer, kind='step'):
    """三种衰减策略，返回 None 表示不衰减。"""
    if kind == 'step':                 # 每 20 轮砍一半
        return optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
    if kind == 'multistep':            # 只在指定轮次砍
        return optim.lr_scheduler.MultiStepLR(
            optimizer, milestones=[50, 125, 160], gamma=0.5)
    if kind == 'exp':                  # 每轮都乘 gamma
        return optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
    return None


# ---------------------------------------------------------------- 插槽三：正则化层
class RegularizedNet(nn.Module):
    """Linear -> BN -> ReLU -> Dropout 是常用的堆叠顺序。"""

    def __init__(self, in_features, n_class, p=0.2):
        super().__init__()
        self.linear1 = nn.Linear(in_features, 128)
        self.bn1 = nn.BatchNorm1d(128)        # 校准站，插在线性层之后
        self.linear2 = nn.Linear(128, 256)
        self.bn2 = nn.BatchNorm1d(256)
        self.dropout = nn.Dropout(p=p)        # 随机休息，插在激活之后
        self.output = nn.Linear(256, n_class)

    def forward(self, x):
        x = torch.relu(self.bn1(self.linear1(x)))
        x = self.dropout(x)
        x = torch.relu(self.bn2(self.linear2(x)))
        x = self.dropout(x)
        # 输出层之前不插 Dropout：随机丢掉最终结论没有任何道理
        return self.output(x)                 # 不加 softmax，交给损失函数


# ---------------------------------------------------------------- 主循环
def train(model, dataset, epochs=5, batch_size=16, lr=1e-3):
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True,
                        drop_last=True)       # 防止最后一批只剩 1 条把 BN 打崩
    criterion = nn.CrossEntropyLoss()
    optimizer = build_optimizer(model, 'adam', lr=lr, weight_decay=1e-4)
    scheduler = build_scheduler(optimizer, 'step')

    for epoch in range(epochs):
        model.train()                         # 打开 Dropout / BN 的训练行为
        total_loss, total_num = 0.0, 0
        for bx, by in loader:
            loss = criterion(model(bx), by)

            optimizer.zero_grad()             # 五步骨架，一个字都没变
            loss.backward()
            optimizer.step()                  # 参数更新：每个 batch 一次

            total_loss += loss.item() * len(by)
            total_num += len(by)

        if scheduler is not None:
            scheduler.step()                  # 学习率衰减：每个 epoch 一次

        # 学习率是否真在按预期衰减，只有打印出来才知道
        print('epoch %3d  loss %.5f  lr %.6f'
              % (epoch + 1, total_loss / total_num,
                 optimizer.param_groups[0]['lr']))
    return model


if __name__ == '__main__':
    torch.manual_seed(0)
    x = torch.randn(200, 20)
    y = torch.randint(0, 4, (200,))
    net = RegularizedNet(in_features=20, n_class=4)
    train(net, TensorDataset(x, y))

"""指数加权平均：Momentum、RMSProp、Adam 的共同零件。

一句话：新值 = 上一次的平均 × β + 当前值 × (1-β)。
β 越大，曲线越平滑，越不理会当前这一下的抖动。
"""
import torch
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

N = 30


def ewa(beta):
    torch.manual_seed(0)
    temperature = torch.randn(size=[N]) * 10     # 造 30 天的随机气温
    avg = []
    for idx, temp in enumerate(temperature, 1):
        if idx == 1:
            avg.append(temp)                      # 第一天没有历史，等于自身
            continue
        avg.append(avg[idx - 2] * beta + (1 - beta) * temp)
    return temperature, avg


if __name__ == '__main__':
    days = torch.arange(1, N + 1)
    for beta in (0.5, 0.9):
        temperature, avg = ewa(beta)
        plt.scatter(days, temperature)            # 散点是原始值
        plt.plot(days, avg, color='r')            # 折线是加权平均
        plt.title('β = %.1f' % beta)
        plt.show()

    # 把同一批数据的两条曲线打印出来对比末值
    print('β=0.5 末值 %.4f' % ewa(0.5)[1][-1])
    print('β=0.9 末值 %.4f  ← 更平缓，更依赖历史' % ewa(0.9)[1][-1])

"""四种优化器在同一个损失函数上各走两步，看更新幅度的差别。

损失统一取 loss = w² / 2，它对 w 的导数就是 w 本身，
所以起点 w=1.0 时第一次的梯度必定是 1.0，四种方法完全可比。
"""
import torch


def run(name, make_optimizer):
    w = torch.tensor([1.0], requires_grad=True)
    optimizer = make_optimizer(w)

    for step in (1, 2):
        loss = ((w ** 2) / 2.0).sum()      # 每步都要重新前向，w 变了
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print('%-10s 第%d次  梯度 %.6f  更新后 w %.6f'
              % (name, step, w.grad.item(), w.item()))
    print()


if __name__ == '__main__':
    # SGD + momentum：带惯性，第二步的位移明显比纯 SGD 大
    run('Momentum', lambda w: torch.optim.SGD([w], lr=0.01, momentum=0.9))

    # AdaGrad：累积平方梯度，学习率只降不升，后期容易走不动
    run('AdaGrad', lambda w: torch.optim.Adagrad([w], lr=0.01))

    # RMSProp：把 AdaGrad 的「求和」换成「指数加权平均」，解决过早熄火
    #          alpha 就是公式里的 β
    run('RMSProp', lambda w: torch.optim.RMSprop([w], lr=0.01, alpha=0.9))

    # Adam：Momentum（一阶矩）+ RMSProp（二阶矩），默认首选
    run('Adam', lambda w: torch.optim.Adam([w], lr=0.01, betas=(0.9, 0.99)))

"""三种学习率衰减策略：调度器每个 epoch 走一步，不是每个 batch。"""
import torch
from torch import optim
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

LR, ITERATION, MAX_EPOCH = 0.1, 10, 200


def collect(make_scheduler, title):
    w = torch.tensor([1.0], requires_grad=True)
    x = torch.tensor([1.0])
    y_true = torch.tensor([0.0])

    optimizer = optim.SGD([w], lr=LR, momentum=0.9)
    scheduler = make_scheduler(optimizer)

    lr_list, epoch_list = [], []
    for epoch in range(MAX_EPOCH):
        lr_list.append(scheduler.get_last_lr()[0])
        epoch_list.append(epoch)

        for _ in range(ITERATION):                 # 一个 epoch 里的每个 batch
            loss = ((w * x - y_true) ** 2) / 2.0
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        scheduler.step()      # 关键：放在 epoch 循环里，不要放进 batch 循环

    plt.plot(epoch_list, lr_list, label=title)
    plt.xlabel('Epoch')
    plt.ylabel('Learning rate')
    plt.legend()
    plt.show()
    print('%s：起始 %.4f -> 结束 %.6f' % (title, lr_list[0], lr_list[-1]))


if __name__ == '__main__':
    # 等间隔：每 50 个 epoch 砍一半
    collect(lambda o: optim.lr_scheduler.StepLR(o, step_size=50, gamma=0.5),
            '等间隔衰减 StepLR')

    # 指定间隔：只在这几个点上砍
    collect(lambda o: optim.lr_scheduler.MultiStepLR(o, milestones=[50, 125, 160],
                                                     gamma=0.5),
            '指定间隔衰减 MultiStepLR')

    # 指数：每个 epoch 都乘 gamma，平滑下降
    collect(lambda o: optim.lr_scheduler.ExponentialLR(o, gamma=0.95),
            '指数衰减 ExponentialLR')

"""完整案例：手机价格区间四分类。

数据：手机价格预测.csv，20 个硬件参数 -> 4 个价格区间（0/1/2/3）。
这一版把本讲的四个优化手段全用上：标准化、Adam、Dropout、更深的网络。
"""
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch.utils.data import TensorDataset, DataLoader

CSV = 'data/手机价格预测.csv'


def create_dataset(csv_path=CSV):
    data = pd.read_csv(csv_path, encoding='UTF-8')

    x = data.iloc[:, :-1]        # 前 20 列是特征
    y = data.iloc[:, -1]         # 最后一列 price_range 是标签

    x_train, x_test, y_train, y_test = train_test_split(
        x, y, test_size=0.2, random_state=10, shuffle=True)

    # 标准化：ram 是几千、clock_speed 是个位数，不消量纲差异，梯度会被大特征带偏
    # 注意 fit 只在训练集上做，测试集只能 transform，否则就是信息泄漏
    transformer = StandardScaler()
    x_train = transformer.fit_transform(x_train)
    x_test = transformer.transform(x_test)

    # 特征必须 float32，分类标签必须 int64，这两个类型错了直接抛 RuntimeError
    train_dataset = TensorDataset(torch.tensor(x_train, dtype=torch.float32),
                                  torch.tensor(y_train.values, dtype=torch.int64))
    test_dataset = TensorDataset(torch.tensor(x_test, dtype=torch.float32),
                                 torch.tensor(y_test.values, dtype=torch.int64))

    return train_dataset, test_dataset, x_train.shape[1], len(np.unique(y_train))


class PhoneANN(nn.Module):
    def __init__(self, features, targets, p=0.2):
        super().__init__()
        self.linear1 = nn.Linear(features, 128)
        self.linear2 = nn.Linear(128, 256)
        self.dropout = nn.Dropout(p=p)        # 只在隐藏层之间插，不要插在输出层前
        self.output = nn.Linear(256, targets)

    def forward(self, x):
        x = torch.relu(self.linear1(x))
        x = self.dropout(x)
        x = torch.relu(self.linear2(x))
        x = self.dropout(x)
        # 这里不写 softmax：CrossEntropyLoss 自带，重复加会把梯度压平
        return self.output(x)


def train_model(train_dataset, features, targets,
                epochs=50, lr=1e-3, batch_size=8, seed=815):
    torch.manual_seed(seed)
    dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    model = PhoneANN(features, targets)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.99))

    for epoch in range(epochs):
        model.train()                       # 打开 Dropout
        total_loss, total_num = 0.0, 0
        for bx, by in dataloader:
            logits = model(bx)
            loss = criterion(logits, by)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * len(by)
            total_num += len(by)
        print('epoch %3d  loss %.5f' % (epoch + 1, total_loss / total_num))

    torch.save(model.state_dict(), 'model/phone_ann.pt')
    return model


def evaluate(model, test_dataset, batch_size=8):
    model.eval()                            # 关闭 Dropout，否则每次预测结果都不同
    dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    correct = 0
    with torch.no_grad():                   # 预测不需要梯度，省显存也更快
        for bx, by in dataloader:
            y_pred = model(bx).argmax(dim=1)
            correct += (y_pred == by).sum().item()

    acc = correct / len(test_dataset)
    print('测试集准确率 %.5f' % acc)
    return acc


def main():
    train_dataset, test_dataset, features, targets = create_dataset()
    print('特征数 %d，类别数 %d' % (features, targets))
    model = train_model(train_dataset, features, targets)
    evaluate(model, test_dataset)


if __name__ == '__main__':
    main()

"""早停（Early Stopping）：验证损失不再下降就收手，并回滚到最好的那一版。

早停是最便宜的正则化手段——不改网络、不改损失函数、不加超参数，
只是「别再练了」。但它有两个新手常漏的细节：

  1. 判据必须是**验证集**指标，拿训练损失判早停等于没判；
  2. 停下来之后要把权重**回滚到最佳轮次**，而不是留着最后一轮的权重——
     最后一轮已经是开始变差的那一版了。

这个类不依赖 torch，可以直接跑；save/restore 两个钩子留给调用方注入。
"""


class EarlyStopping:
    def __init__(self, patience=5, min_delta=1e-4, mode='min'):
        """patience : 连续多少轮没有改善就停
        min_delta: 改善多少才算数，防止在噪声里反复横跳
        mode     : 'min' 盯损失越小越好，'max' 盯准确率越大越好
        """
        self.patience = patience
        self.min_delta = min_delta
        self.mode = mode
        self.best = None
        self.best_epoch = -1
        self.bad_epochs = 0
        self.should_stop = False

    def _is_better(self, value):
        if self.best is None:
            return True
        if self.mode == 'min':
            return value < self.best - self.min_delta
        return value > self.best + self.min_delta

    def step(self, value, epoch, on_improve=None):
        """每个 epoch 算完验证指标后调用一次。返回是否应当停止。"""
        if self._is_better(value):
            self.best = value
            self.best_epoch = epoch
            self.bad_epochs = 0
            if on_improve is not None:
                on_improve()               # 这里把当前权重存下来
        else:
            self.bad_epochs += 1
            if self.bad_epochs >= self.patience:
                self.should_stop = True
        return self.should_stop


def train_with_early_stopping(model, train_loader, valid_loader,
                              criterion, optimizer, max_epochs=200,
                              patience=5, ckpt='model/best.pt'):
    """接进真实训练循环的样子。"""
    import copy
    import torch

    stopper = EarlyStopping(patience=patience, mode='min')
    best_state = None

    for epoch in range(max_epochs):
        # ---- 训练一轮
        model.train()                       # 打开 Dropout / BN 的训练行为
        for bx, by in train_loader:
            loss = criterion(model(bx), by)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # ---- 在验证集上评估
        model.eval()                        # 关掉 Dropout，BN 改用全局统计量
        total, num = 0.0, 0
        with torch.no_grad():               # 评估不需要梯度
            for bx, by in valid_loader:
                total += criterion(model(bx), by).item() * len(by)
                num += len(by)
        valid_loss = total / num
        print('epoch %3d  valid_loss %.5f' % (epoch + 1, valid_loss))

        def keep():
            nonlocal best_state
            best_state = copy.deepcopy(model.state_dict())

        if stopper.step(valid_loss, epoch, on_improve=keep):
            print('验证损失连续 %d 轮没有改善，停在第 %d 轮' % (patience, epoch + 1))
            break

    # ---- 回滚到最好的那一版，这一步不能省
    if best_state is not None:
        model.load_state_dict(best_state)
        torch.save(best_state, ckpt)
        print('已回滚到第 %d 轮的权重，验证损失 %.5f'
              % (stopper.best_epoch + 1, stopper.best))
    return model


if __name__ == '__main__':
    # 用 overfit_diagnose.py 里那条过拟合曲线走一遍，看早停停在哪
    valid_curve = [1.42, 1.00, 0.72, 0.61, 0.58, 0.60, 0.66, 0.74, 0.83, 0.91]
    stopper = EarlyStopping(patience=3, mode='min')
    for epoch, v in enumerate(valid_curve):
        stopped = stopper.step(v, epoch)
        print('epoch %2d  valid %.2f  best %.2f(第%d轮)  连续变差 %d 轮  %s'
              % (epoch + 1, v, stopper.best, stopper.best_epoch + 1,
                 stopper.bad_epochs, '-> 停止' if stopped else ''))
        if stopped:
            break

"""调优骨架模板：把可调的旋钮集中到一处，一次只动一个。"""
import torch
from torch import nn, optim
from torch.utils.data import DataLoader

CONFIG = {
    'lr': 1e-3,            # TODO: 学不动就调大，损失来回震荡就调小
    'batch_size': 8,       # TODO: 显存够就调大，大 batch 更稳但更新次数少
    'epochs': 50,          # TODO: 看损失曲线，压平了就可以停
    'dropout': 0.2,        # TODO: 训练集好测试集差（过拟合）就调大
    'weight_decay': 0.0,   # TODO: 需要 L2 正则时填 1e-4 量级
    'optimizer': 'adam',   # TODO: 'adam' / 'sgd'
    'scheduler': 'step',   # TODO: 'none' / 'step' / 'exp'
    'seed': 0,
}


def build_optimizer(model, cfg=CONFIG):
    if cfg['optimizer'] == 'adam':
        return optim.Adam(model.parameters(), lr=cfg['lr'],
                          weight_decay=cfg['weight_decay'])
    return optim.SGD(model.parameters(), lr=cfg['lr'], momentum=0.9,
                     weight_decay=cfg['weight_decay'])


def build_scheduler(optimizer, cfg=CONFIG):
    if cfg['scheduler'] == 'step':
        return optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
    if cfg['scheduler'] == 'exp':
        return optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
    return None


def train(model, dataset, cfg=CONFIG):
    torch.manual_seed(cfg['seed'])
    loader = DataLoader(dataset, batch_size=cfg['batch_size'], shuffle=True)

    criterion = nn.CrossEntropyLoss()      # TODO: 回归换成 MSELoss / SmoothL1Loss
    optimizer = build_optimizer(model, cfg)
    scheduler = build_scheduler(optimizer, cfg)

    for epoch in range(cfg['epochs']):
        model.train()
        total_loss, total_num = 0.0, 0
        for bx, by in loader:
            loss = criterion(model(bx), by)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * len(by)
            total_num += len(by)

        if scheduler is not None:
            scheduler.step()               # 每个 epoch 一次，别放进内层循环

        print('epoch %3d  loss %.5f  lr %.6f'
              % (epoch + 1, total_loss / total_num,
                 optimizer.param_groups[0]['lr']))
    return model