"""优化器：AdamW / 学习率调度""" # 2026-01-23 import math import torch.optim as optim from torch.optim.lr_scheduler import LambdaLR def get_optimizer(model, config): """ 获取优化器（AdamW） 参数: model: 模型 config: 训练配置字典，应包含以下键： - learning_rate: 学习率（默认: 1e-4） - weight_decay: 权重衰减（默认: 0.01） - beta1: Adam beta1（默认: 0.9） - beta2: Adam beta2（默认: 0.999） - eps: Adam epsilon（默认: 1e-8） 返回: AdamW 优化器 """ lr = float(config.get("learning_rate", 1e-4)) weight_decay = float(config.get("weight_decay", 0.01)) beta1 = float(config.get("beta1", 0.9)) beta2 = float(config.get("beta2", 0.999)) eps = float(config.get("eps", 1e-8)) optimizer = optim.AdamW( model.parameters(), lr=lr, weight_decay=weight_decay, betas=(beta1, beta2), eps=eps, ) return optimizer def get_lr_scheduler(optimizer, config): """ 获取学习率调度器（支持预热） 参数: optimizer: 优化器 config: 训练配置字典，应包含以下键： - lr_scheduler: 调度器类型，可选值： - "cosine": 余弦退火（带预热） - "linear": 线性衰减（带预热） - "constant": 常数学习率 - warmup_steps: 预热步数（默认: 100） - max_steps: 最大训练步数（默认: 10000） 返回: 学习率调度器（如果为 constant，返回 None） """ scheduler_type = config.get("lr_scheduler", "cosine") max_steps = config.get("max_steps", 10000) warmup_steps = config.get("warmup_steps", 100) if scheduler_type == "cosine": # 余弦退火调度（带预热） def lr_lambda(step): if step < warmup_steps: # 预热阶段：线性增加从 0 到 1 return step / warmup_steps if warmup_steps > 0 else 1.0 else: # 余弦退火阶段：从 1 衰减到 0 # 处理边界情况：如果 max_steps <= warmup_steps，直接返回最小值 if max_steps <= warmup_steps: return 0.0 progress = (step - warmup_steps) / (max_steps - warmup_steps) # 限制 progress 在 [0, 1] 范围内 progress = min(progress, 1.0) # 余弦退火：0.5 * (1 + cos(π * progress)) return 0.5 * (1.0 + math.cos(progress * math.pi)) scheduler = LambdaLR(optimizer, lr_lambda) elif scheduler_type == "linear": # 线性衰减调度（带预热） def lr_lambda(step): if step < warmup_steps: # 预热阶段：线性增加从 0 到 1 return step / warmup_steps if warmup_steps > 0 else 1.0 else: # 线性衰减阶段：从 1 线性衰减到 0.1 # 处理边界情况：如果 max_steps <= warmup_steps，直接返回最小值 if max_steps <= warmup_steps: return 0.1 progress = (step - warmup_steps) / (max_steps - warmup_steps) # 限制 progress 在 [0, 1] 范围内 progress = min(progress, 1.0) # 线性衰减：从 1.0 到 0.1 return 1.0 - 0.9 * progress scheduler = LambdaLR(optimizer, lr_lambda) elif scheduler_type == "constant": # 常数学习率（无调度器） scheduler = None else: raise ValueError( f"未知的学习率调度器类型: {scheduler_type}。" f"支持的类型: cosine, linear, constant" ) return scheduler if __name__ == "__main__": import sys import io # 设置输出编码为 UTF-8（Windows 兼容） if sys.platform == "win32": sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8") print("=" * 60) print("优化器和学习率调度器测试") print("=" * 60) # 创建一个简单的模型用于测试 import torch.nn as nn class SimpleModel(nn.Module): def __init__(self): super().__init__() self.linear = nn.Linear(10, 1) def forward(self, x): return self.linear(x) model = SimpleModel() # 测试配置 config = { "learning_rate": 1e-3, "weight_decay": 0.01, "beta1": 0.9, "beta2": 0.999, "eps": 1e-8, "lr_scheduler": "cosine", "warmup_steps": 10, "max_steps": 100, } # 测试优化器 print("\n1. 测试优化器创建") optimizer = get_optimizer(model, config) print(f" 优化器类型: {type(optimizer).__name__}") print(f" 学习率: {optimizer.param_groups[0]['lr']}") print(f" 权重衰减: {optimizer.param_groups[0]['weight_decay']}") print(f" Beta1: {optimizer.param_groups[0]['betas'][0]}") print(f" Beta2: {optimizer.param_groups[0]['betas'][1]}") # 测试学习率调度器（cosine） print("\n2. 测试余弦退火学习率调度器（带预热）") scheduler = get_lr_scheduler(optimizer, config) print(f" 调度器类型: {type(scheduler).__name__}") print(f" 初始学习率: {optimizer.param_groups[0]['lr']:.6f}") # 模拟训练步骤，观察学习率变化 print("\n3. 模拟训练步骤，观察学习率变化（cosine）") lrs = [] for step in range(0, 101, 10): current_lr = optimizer.param_groups[0]["lr"] lrs.append(current_lr) print(f" 步数 {step:3d}: 学习率 = {current_lr:.6f}") # 模拟优化步骤（先 optimizer.step()，再 scheduler.step()） # 创建一个虚拟的损失并反向传播（仅用于测试） dummy_loss = sum(p.sum() for p in model.parameters()) dummy_loss.backward() optimizer.step() optimizer.zero_grad() if scheduler is not None: scheduler.step() # 测试线性调度器 print("\n4. 测试线性学习率调度器（带预热）") config_linear = config.copy() config_linear["lr_scheduler"] = "linear" optimizer_linear = get_optimizer(model, config_linear) scheduler_linear = get_lr_scheduler(optimizer_linear, config_linear) print(f" 调度器类型: {type(scheduler_linear).__name__}") print("\n5. 模拟训练步骤，观察学习率变化（linear）") for step in range(0, 101, 10): current_lr = optimizer_linear.param_groups[0]["lr"] print(f" 步数 {step:3d}: 学习率 = {current_lr:.6f}") # 模拟优化步骤（先 optimizer.step()，再 scheduler.step()） dummy_loss = sum(p.sum() for p in model.parameters()) dummy_loss.backward() optimizer_linear.step() optimizer_linear.zero_grad() if scheduler_linear is not None: scheduler_linear.step() # 测试常数学习率 print("\n6. 测试常数学习率") config_constant = config.copy() config_constant["lr_scheduler"] = "constant" optimizer_constant = get_optimizer(model, config_constant) scheduler_constant = get_lr_scheduler(optimizer_constant, config_constant) print(f" 调度器类型: {scheduler_constant}") print(f" 学习率: {optimizer_constant.param_groups[0]['lr']:.6f}") # 测试错误类型 print("\n7. 测试错误类型处理") try: config_error = config.copy() config_error["lr_scheduler"] = "invalid" scheduler_error = get_lr_scheduler(optimizer, config_error) except ValueError as e: print(f" 正确捕获错误: {e}") print("\n" + "=" * 60) print("所有测试完成！") print("=" * 60)