第 6 章：优化器与损失函数¶

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

# 创建一个简单的二次优化问题
# 目标：找到 f(x, y) = x^2 + 5*y^2 的最小值（显然在 (0, 0)）

def create_optimization_problem():
    """创建优化问题：最小化 f(x, y) = x^2 + 5*y^2"""
    x = torch.tensor([3.0, 3.0], requires_grad=True)
    return x

def optimize_and_track(optimizer_class, lr, steps=50, **kwargs):
    """运行优化并记录轨迹"""
    x = create_optimization_problem()
    optimizer = optimizer_class([x], lr=lr, **kwargs)
    trajectory = [x.detach().clone().numpy()]

    for _ in range(steps):
        optimizer.zero_grad()
        loss = x[0]**2 + 5 * x[1]**2
        loss.backward()
        optimizer.step()
        trajectory.append(x.detach().clone().numpy())

    return np.array(trajectory)

# 对比不同优化器
optimizers_config = {
    'SGD (lr=0.1)': (optim.SGD, 0.1),
    'SGD+Momentum (lr=0.1)': (optim.SGD, 0.1, {'momentum': 0.9}),
    'Adam (lr=0.1)': (optim.Adam, 0.1),
    'RMSprop (lr=0.1)': (optim.RMSprop, 0.1),
}

trajectories = {}
for name, (opt_class, lr, *args) in optimizers_config.items():
    kwargs = args[0] if args else {}
    trajectories[name] = optimize_and_track(opt_class, lr, **kwargs)

# 可视化
fig, axes = plt.subplots(1, 4, figsize=(16, 4))
for ax, (name, traj) in zip(axes, trajectories.items()):
    ax.plot(traj[:, 0], traj[:, 1], 'b-o', markersize=3)
    ax.plot(0, 0, 'r*', markersize=10, label='最优解')
    ax.set_xlim(-1, 4)
    ax.set_ylim(-1, 4)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title(name)
    ax.legend()
    ax.grid(True)

plt.suptitle('不同优化器的优化轨迹对比\nf(x,y) = x² + 5y², 起点 (3,3)')
plt.tight_layout()
plt.show()

# 打印最终结果
print("各优化器最终结果：")
for name, traj in trajectories.items():
    final = traj[-1]
    loss = final[0]**2 + 5 * final[1]**2
    print(f"  {name:<25} x=({final[0]:.4f}, {final[1]:.4f}), loss={loss:.6f}")

def test_learning_rates(lrs):
    """测试不同学习率对 SGD 的影响"""
    results = {}

    for lr in lrs:
        x = create_optimization_problem()
        optimizer = optim.SGD([x], lr=lr)
        losses = []

        for _ in range(30):
            optimizer.zero_grad()
            loss = x[0]**2 + 5 * x[1]**2
            loss.backward()
            optimizer.step()
            losses.append(loss.item())

        results[lr] = losses

    return results

lrs = [0.01, 0.05, 0.1, 0.2, 0.5]
results = test_learning_rates(lrs)

plt.figure(figsize=(10, 5))
for lr, losses in results.items():
    plt.plot(losses, label=f'lr={lr}')

plt.xlabel('迭代步数')
plt.ylabel('损失值')
plt.title('不同学习率下 SGD 的收敛曲线')
plt.legend()
plt.grid(True)
plt.yscale('log')
plt.show()

print("最终损失值：")
for lr, losses in results.items():
    print(f"  lr={lr:.2f}: {losses[-1]:.6f}")

# 创建模型和优化器
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 1)
)
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 不同调度策略
schedulers = {}

# StepLR：每 step_size 个 epoch 乘以 gamma
schedulers['StepLR'] = optim.lr_scheduler.StepLR(
    optimizer, step_size=30, gamma=0.1
)

# MultiStepLR：在指定 epoch 乘以 gamma
schedulers['MultiStepLR'] = optim.lr_scheduler.MultiStepLR(
    optimizer, milestones=[30, 80], gamma=0.1
)

# ExponentialLR：每个 epoch 乘以 gamma
schedulers['ExponentialLR'] = optim.lr_scheduler.ExponentialLR(
    optimizer, gamma=0.95
)

# CosineAnnealingLR：余弦退火
schedulers['CosineAnnealingLR'] = optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=100
)

# 可视化各调度策略的学习率变化
epochs = 100
lr_records = {}

for name, scheduler in schedulers.items():
    # 重置优化器学习率
    for param_group in optimizer.param_groups:
        param_group['lr'] = 0.1

    lrs = []
    for epoch in range(epochs):
        lrs.append(optimizer.param_groups[0]['lr'])
        scheduler.step()
    lr_records[name] = lrs

plt.figure(figsize=(12, 5))
for name, lrs in lr_records.items():
    plt.plot(lrs, label=name)

plt.xlabel('Epoch')
plt.ylabel('Learning Rate')
plt.title('不同学习率调度策略对比')
plt.legend()
plt.grid(True)
plt.show()

class WeightedMSELoss(nn.Module):
    """带权重的 MSE 损失：对某些样本给予更高权重"""
    def __init__(self, weight_positive=2.0):
        super(WeightedMSELoss, self).__init__()
        self.weight_positive = weight_positive

    def forward(self, pred, target):
        # 基础 MSE
        mse = (pred - target) ** 2

        # 对正样本（target > 0）给予更高权重
        weights = torch.where(target > 0, self.weight_positive, 1.0)

        return (mse * weights).mean()

# 测试自定义损失
pred = torch.tensor([0.5, 1.5, 2.0, 0.8])
target = torch.tensor([0.0, 2.0, 2.0, 1.0])

standard_mse = nn.MSELoss()
weighted_mse = WeightedMSELoss(weight_positive=3.0)

print(f"标准 MSE:     {standard_mse(pred, target).item():.4f}")
print(f"加权 MSE (x3): {weighted_mse(pred, target).item():.4f}")

# 正确示例：多分类
logits = torch.randn(4, 10)  # 模型原始输出，没有 Softmax
targets = torch.randint(0, 10, (4,))

# CrossEntropyLoss 内部包含 Softmax
criterion = nn.CrossEntropyLoss()
loss = criterion(logits, targets)
print(f"CrossEntropyLoss(logits): {loss.item():.4f}")

# 错误示例：手动加 Softmax 再用 CrossEntropyLoss
probs = torch.softmax(logits, dim=1)
# loss_wrong = criterion(probs, targets)  # 这会导致双重 Softmax！

model = nn.Linear(10, 1)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 模拟训练几个 step
for _ in range(10):
    x = torch.randn(4, 10)
    y = torch.randn(4, 1)

    optimizer.zero_grad()
    loss = nn.MSELoss()(model(x), y)
    loss.backward()
    optimizer.step()

# 保存检查点
checkpoint = {
    'epoch': 10,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss.item(),
}
torch.save(checkpoint, 'checkpoint.pth')
print(f"检查点已保存: epoch={checkpoint['epoch']}, loss={checkpoint['loss']:.4f}")

# 恢复训练
new_model = nn.Linear(10, 1)
new_optimizer = optim.Adam(new_model.parameters(), lr=0.001)

checkpoint = torch.load('checkpoint.pth')
new_model.load_state_dict(checkpoint['model_state_dict'])
new_optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']

print(f"从 epoch {start_epoch} 恢复训练")

import os
os.remove('checkpoint.pth')