import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import numpy as np

# 生成合成数据：y = 2x + 1 + 噪声
np.random.seed(42)
X_numpy = np.random.rand(100, 1) * 10
y_numpy = 2 * X_numpy + 1 + np.random.randn(100, 1) * 2 # 添加一些噪声

# 转换为PyTorch张量
X = torch.tensor(X_numpy, dtype=torch.float32)
y = torch.tensor(y_numpy, dtype=torch.float32)

# 定义一个简单的线性模型
model = nn.Linear(1, 1)

# 定义损失函数（均方误差）
criterion = nn.MSELoss()

# 创建数据集
dataset = TensorDataset(X, y)

我们有100个数据点，其中 $y$ 大约等于 $2x + 1$ 。我们的 nn.Linear(1, 1) 模型尝试学习权重（斜率，目标值为2）和偏置（截距，目标值为1）。我们将使用均方误差 (MSE) 作为损失函数。

定义优化器

我们将比较三种优化方法：

小批量梯度下降（批量大小16）： 一种常用的实用方法。我们将使用PyTorch的 SGD 优化器，搭配标准学习率。
随机梯度下降（批量大小1）： 小批量的极端情况，在处理每个样本后进行更新。
带动量的SGD（批量大小16）： 使用动量变体来可能加速收敛。

每个实验我们需要独立的模型实例和优化器，以确保公平比较。

# 学习率
learning_rate = 0.001
momentum_factor = 0.9
num_epochs = 50

# --- 优化器 ---
# 1. 小批量梯度下降（批量大小16）
model_minibatch = nn.Linear(1, 1)
optimizer_minibatch = optim.SGD(model_minibatch.parameters(), lr=learning_rate)
dataloader_minibatch = DataLoader(dataset, batch_size=16, shuffle=True)

# 2. 随机梯度下降（批量大小1）
model_sgd = nn.Linear(1, 1)
optimizer_sgd = optim.SGD(model_sgd.parameters(), lr=learning_rate)
dataloader_sgd = DataLoader(dataset, batch_size=1, shuffle=True)

# 3. 带动量的SGD（批量大小16）
model_momentum = nn.Linear(1, 1)
optimizer_momentum = optim.SGD(model_momentum.parameters(), lr=learning_rate, momentum=momentum_factor)
# 为了公平比较，我们使用与小批量梯度下降相同的数据加载器
dataloader_momentum = DataLoader(dataset, batch_size=16, shuffle=True)

注意，对于纯粹的SGD，我们将 batch_size 设置为1。对于小批量和动量法，我们使用 batch_size=16。动量优化器与小批量优化器相同，只是多了一个 momentum=0.9 参数。

训练循环

所有优化器的训练循环结构类似。我们遍历各个周期（epoch），在每个周期内，我们遍历由 DataLoader 提供的数据批次。

def train_model(model, optimizer, dataloader, criterion, epochs):
    """辅助函数，用于训练模型并记录每个周期的损失。"""
    epoch_losses = []
    for epoch in range(epochs):
        epoch_loss = 0.0
        num_batches = 0
        for inputs, targets in dataloader:
            # 将参数梯度归零
            optimizer.zero_grad()

            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            # 反向传播并优化
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()
            num_batches += 1

        avg_epoch_loss = epoch_loss / num_batches
        epoch_losses.append(avg_epoch_loss)
        # 可选：打印进度
        # if (epoch + 1) % 10 == 0:
        #     print(f'Epoch [{epoch+1}/{epochs}], Loss: {avg_epoch_loss:.4f}')
    return epoch_losses

# 训练每个模型
losses_minibatch = train_model(model_minibatch, optimizer_minibatch, dataloader_minibatch, criterion, num_epochs)
losses_sgd = train_model(model_sgd, optimizer_sgd, dataloader_sgd, criterion, num_epochs)
losses_momentum = train_model(model_momentum, optimizer_momentum, dataloader_momentum, criterion, num_epochs)

print("训练完成。")

比较性能

比较这些优化器最直接的方法是绘制它们在不同周期（epoch）上的训练损失。损失越低通常表示模型拟合越好，下降越快则表示收敛越迅速。

在合成线性回归任务中，小批量梯度下降（批量大小16）、随机梯度下降（批量大小1）和带动量的SGD（批量大小16）的每个周期的平均训练损失。请注意y轴采用对数刻度。

结果分析

让我们分析这张图表（请记住，由于随机初始化和数据混洗，结果可能略有不同）：

小批量梯度下降（蓝线）： 损失呈现相对平滑的下降。使用16个样本的批次平均了纯粹SGD中存在的部分梯度噪声，从而实现稳定的收敛。它能相对较快地达到较低的损失值。
SGD（橙线）： 损失曲线明显噪声更大，尤其是在初始周期。由于更新基于单个样本，梯度波动可能很大，导致损失值跳动更多。虽然它最终会收敛，但在本例中，它需要比小批量梯度下降或动量法更多的周期才能达到相似的损失水平。噪声有时有助于跳出不理想的局部最小值，但通常会减慢收敛速度。
动量法（绿线）： 这通常表现出最快的初始收敛速度。动量项有助于在一致的下降方向上加速进展，并抑制纯SGD或小批量梯度下降可能出现的振荡。它能迅速稳定在最低损失值附近。在这个简单案例中，它的收敛速度明显快于其他两种方法。

总结

这种实践比较突出了我们之前讨论的特点：

批量大小很重要： 小批量梯度下降（大小>1）在SGD的计算效率和批量梯度下降的稳定收敛之间提供了一个平衡（由于它在大数据集上的效率低下，我们没有运行）。它的收敛比纯SGD更平滑。
SGD有噪声： 使用批量大小为1会在梯度估计中引入显著的方差，导致优化路径噪声较大。这有时可能带来益处，但通常会减慢收敛速度。
动量法加速： 通过引入速度项，动量法通常比纯SGD或小批量梯度下降收敛更快，特别是在梯度方向相对一致或在浅层区域移动时。

这个练习表明，即使是核心算法也有其不同的表现。尽管小批量SGD是一种主力方法，动量法通常能提供显著的加速，为我们将在下一章中考察的自适应方法铺平道路。请记住，最佳选择通常取决于具体问题、数据集和模型架构。实验是经常需要的。

这部分内容有帮助吗？