# Fixed learning rate SGD
lr = 0.01  # 全局固定学习率
for epoch in range(num_epochs):
    for x, y in dataloader:
        loss = model(x, y)
        loss.backward()
        for param in model.parameters():
            param.data -= lr * param.grad

# Step decay
def get_lr(epoch, init_lr=0.01, gamma=0.1, step_size=30):
    return init_lr * (gamma ** (epoch // step_size))

# Cosine annealing
import math
def cosine_lr(epoch, T, eta_min=1e-5, eta_max=0.01):
    return eta_min + 0.5 * (eta_max - eta_min) * (1 + math.cos(math.pi * epoch / T))

# Simplified Adam optimizer
class Adam:
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999)):
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.m = [torch.zeros_like(p) for p in params]  # 1st moment
        self.v = [torch.zeros_like(p) for p in params]  # 2nd moment

    def step(self, params, grads, t):
        for i, (param, grad) in enumerate(zip(params, grads)):
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * grad**2
            m_hat = self.m[i] / (1 - self.beta1**t)   # bias correction
            v_hat = self.v[i] / (1 - self.beta2**t)   # bias correction
            param.data -= self.lr * m_hat / (v_hat.sqrt() + 1e-8)

# Discriminative fine-tuning in PyTorch
import torch

def get_discriminative_lrs(model, base_lr=0.01, decay_factor=2.6):
    """
    Compute per-layer learning rates for discriminative fine-tuning.

    Args:
        model: a model with L layers
        base_lr: learning rate for the LAST layer
        decay_factor: each lower layer's LR = upper layer's LR / decay_factor

    Returns:
        list of (param_group, lr) tuples
    """
    layers = list(model.children())  # e.g., [layer1, layer2, layer3]
    num_layers = len(layers)

    param_groups = []
    for i, layer in enumerate(layers):
        # Layer index: 0 = bottom (smallest LR), L-1 = top (largest LR)
        lr = base_lr / (decay_factor ** (num_layers - 1 - i))
        param_groups.append({
            'params': layer.parameters(),
            'lr': lr
        })

    return param_groups

# Usage example
model = ThreeLayerLSTM()  # AWD-LSTM with 3 layers
param_groups = get_discriminative_lrs(model, base_lr=0.01, decay_factor=2.6)

# Layer 1 (bottom): lr ≈ 0.00148
# Layer 2 (middle): lr ≈ 0.00385
# Layer 3 (top):    lr = 0.01

optimizer = torch.optim.Adam(param_groups)

import torch.nn as nn
import torch.optim as optim

class ULMPfinetuner:
    def __init__(self, model, base_lr=0.01, decay_factor=2.6):
        self.model = model
        self.base_lr = base_lr
        self.decay_factor = decay_factor

    def get_param_groups(self):
        """
        Build parameter groups with discriminative learning rates.

        Architecture (bottom to top):
          - Embedding layer:   smallest LR (most general)
          - RNN layer 1:
          - RNN layer 2:
          - RNN layer 3:
          - Classifier head:   largest LR (most task-specific)
        """
        layer_names = [
            'embedding',       # most general
            'rnn_layer_0',
            'rnn_layer_1',
            'rnn_layer_2',
            'classifier',      # most task-specific
        ]

        param_groups = []
        for i, name in enumerate(layer_names):
            lr = self.base_lr / (self.decay_factor ** (len(layer_names) - 1 - i))
            module = getattr(self.model, name)
            param_groups.append({
                'params': list(module.parameters()),
                'lr': lr,
                'name': name,
            })
            print(f"  {name:20s} lr = {lr:.6f}")

        return param_groups

# Instantiate and train
finetuner = ULMPfinetuner(model, base_lr=0.01, decay_factor=2.6)
optimizer = optim.Adam(finetuner.get_param_groups())

for epoch in range(num_epochs):
    for batch in dataloader:
        loss = compute_loss(model, batch)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

import numpy as np
import matplotlib.pyplot as plt

def slanted_triangular_lr(t, T, eta_max=0.01, cut_frac=0.1, ratio=32):
    """
    Slanted Triangular Learning Rate (STLR) schedule.

    Phase 1 (t < cut): LR linearly increases from eta_max/ratio to eta_max
    Phase 2 (t >= cut): LR linearly decreases from eta_max to eta_max/ratio

    Args:
        t: current iteration
        T: total number of training iterations
        eta_max: maximum learning rate
        cut_frac: fraction of iterations to increase LR
        ratio: how much smaller the min LR is vs max LR
    """
    cut = int(T * cut_frac)
    if cut == 0:
        cut = 1

    if t < cut:
        p = t / cut
    else:
        p = 1.0 - (t - cut) / (cut * (1.0 / cut_frac - 1.0))

    return eta_max * (1.0 + p * (ratio - 1.0)) / ratio

# Visualization
T = 10000
iterations = np.arange(T)
lrs = [slanted_triangular_lr(t, T) for t in iterations]

plt.figure(figsize=(10, 4))
plt.plot(iterations, lrs, linewidth=2, color='#2196F3')
plt.xlabel('Training Iteration', fontsize=12)
plt.ylabel('Learning Rate', fontsize=12)
plt.title('Slanted Triangular Learning Rate', fontsize=14)
plt.axvline(x=int(T*0.1), color='red', linestyle='--', alpha=0.5, label='cut point')
plt.legend(fontsize=11)
plt.tight_layout()
plt.show()

# Full ULMFiT learning rate: Discriminative + STLR
class ULMLearningRateScheduler:
    def __init__(self, num_layers, T, base_lr_max=0.01,
                 decay_factor=2.6, cut_frac=0.1, ratio=32):
        self.T = T
        self.cut_frac = cut_frac
        self.ratio = ratio

        # Compute per-layer max learning rates (discriminative)
        self.lr_max_per_layer = []
        for layer_idx in range(num_layers):
            lr_max = base_lr_max / (decay_factor ** (num_layers - 1 - layer_idx))
            self.lr_max_per_layer.append(lr_max)

    def get_lr(self, t, layer_idx):
        """Get learning rate for layer `layer_idx` at iteration `t`."""
        eta_max = self.lr_max_per_layer[layer_idx]
        return slanted_triangular_lr(
            t, self.T,
            eta_max=eta_max,
            cut_frac=self.cut_frac,
            ratio=self.ratio
        )

# Example: 3-layer model
scheduler = ULMLearningRateScheduler(
    num_layers=3,
    T=10000,
    base_lr_max=0.01,     # top layer max LR
    decay_factor=2.6,
)

# Print the learning rate ranges
for i in range(3):
    lr_range = (scheduler.get_lr(1000, i), scheduler.get_lr(0, i))
    print(f"Layer {i+1}: LR range = [{min(lr_range):.6f}, {max(lr_range):.6f}]")

# Output:
# Layer 1 (bottom): LR range ≈ [0.000046, 0.00148]  ← smallest
# Layer 2 (middle): LR range ≈ [0.000120, 0.00385]
# Layer 3 (top):    LR range ≈ [0.000313, 0.01000]   ← largest