Python实现模型参数的平滑（Averaging）技术：SWA/EMA在训练末期的应用

Python实现模型参数的平滑（Averaging）技术：SWA/EMA在训练末期的应用

大家好！今天我们要深入探讨一种在深度学习模型训练中非常有效的技术：模型参数的平滑，特别是其在训练末期的应用。我们将重点关注两种主流的平滑方法：随机权重平均（Stochastic Weight Averaging, SWA）和指数移动平均（Exponential Moving Average, EMA），并提供详细的Python实现代码。

1. 背景与动机：为何需要模型参数平滑？

在深度学习模型的训练过程中，损失函数通常是一个高度非凸的函数。这意味着训练过程会存在很多局部最小值。标准的随机梯度下降（SGD）及其变体（如Adam）在训练过程中可能会陷入这些局部最小值，导致模型的泛化能力受到限制。

模型参数平滑的核心思想是，通过对训练过程中获得的多个模型参数进行平均，来得到一个更鲁棒、泛化能力更强的模型。这种平均可以有效地“平滑”损失函数曲面，使模型参数更接近全局最优解，或者至少位于一个更宽广的局部最优区域，从而提高模型的泛化能力和鲁棒性。

2. 随机权重平均（SWA）：一种简单而有效的平均策略

2.1 SWA的核心思想

SWA是一种简单的平均策略，它在训练周期的后期，以一定的频率保存模型的权重，并在训练结束后对这些权重进行平均。与传统的模型集成方法相比，SWA只需要训练一次模型，大大降低了计算成本。

2.2 SWA的算法流程

1. 正常训练： 使用标准的优化算法（如SGD、Adam）训练模型一段时间。
2. 启动SWA： 在训练的后期，设定一个开始SWA的 epoch swa_start。
3. 权重收集： 从 swa_start 开始，每隔一定的步长（例如每个epoch），保存当前模型的权重。
4. 权重平均： 在训练结束后，对所有保存的权重进行平均，得到最终的SWA模型。

2.3 SWA的Python实现

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

# 假设我们有一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

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

# 假设我们有一个自定义数据集
class DummyDataset(Dataset):
    def __init__(self, num_samples, input_size):
        self.num_samples = num_samples
        self.input_size = input_size
        self.data = torch.randn(num_samples, input_size)
        self.labels = torch.randint(0, 2, (num_samples,))  # 二分类问题

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

def swa_implementation(model, train_loader, optimizer, swa_start, swa_lr, swa_c_epochs, device):
    """
    SWA的实现函数。

    Args:
        model: PyTorch模型。
        train_loader: 训练数据加载器。
        optimizer: 优化器。
        swa_start: 开始SWA的epoch。
        swa_lr: SWA的学习率。
        swa_c_epochs: SWA循环周期的epoch数。
        device: 设备 (CPU 或 CUDA)。
    """

    swa_model = type(model)(*model.args, **model.kwargs).to(device)  # 创建SWA模型副本
    swa_n = 0  # SWA模型平均的次数

    def update_swa():
        """更新SWA模型的权重"""
        nonlocal swa_n
        swa_n += 1
        for name, param in model.named_parameters():
            swa_param = next(swa_model.named_parameters())[1]
            swa_param.data.copy_(
                (swa_param.data * (swa_n - 1) + param.data) / swa_n
            )

    # SWA学习率调度器
    swa_scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=swa_c_epochs, eta_min=swa_lr
    )

    # 训练循环
    for epoch in range(num_epochs):
        for i, (inputs, labels) in enumerate(train_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)

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

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

        # 在swa_start之后启动SWA
        if epoch >= swa_start:
            update_swa() # 更新SWA模型

            # 使用SWA学习率调度器调整学习率
            swa_scheduler.step()

        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}, LR: {optimizer.param_groups[0]['lr']:.6f}")

    # 返回SWA模型
    return swa_model

# 超参数设置
input_size = 10
hidden_size = 20
output_size = 2
num_epochs = 20
batch_size = 32
learning_rate = 0.001
swa_start = 10
swa_lr = 0.0005
swa_c_epochs = 5

# 创建模型、数据集和数据加载器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleModel(input_size, hidden_size, output_size)
model.args = (input_size, hidden_size, output_size) # 保存模型参数，用于创建swa_model副本
model.kwargs = {}

model = model.to(device)
dataset = DummyDataset(num_samples=1000, input_size=input_size)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 执行SWA
swa_model = swa_implementation(model, train_loader, optimizer, swa_start, swa_lr, swa_c_epochs, device)

# 使用SWA模型进行推理 (示例)
def evaluate_model(model, data_loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

test_dataset = DummyDataset(num_samples=200, input_size=input_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 评估原始模型
accuracy_original = evaluate_model(model, test_loader, device)
print(f"原始模型准确率: {accuracy_original:.4f}")

# 评估SWA模型
accuracy_swa = evaluate_model(swa_model, test_loader, device)
print(f"SWA模型准确率: {accuracy_swa:.4f}")

代码解释：

• SimpleModel：一个简单的线性模型，用于演示SWA。
• DummyDataset：一个模拟数据集，用于训练模型。
• swa_implementation：SWA的核心实现函数。它接收模型、数据加载器、优化器等参数，并在训练过程中应用SWA。
• update_swa：用于更新SWA模型的权重，将当前模型的权重累加到SWA模型中，并计算平均值。
• swa_scheduler: 使用余弦退火学习率，SWA论文中推荐使用循环退火学习率，这里使用了简化的余弦退火，更加方便。
• evaluate_model：评估模型在测试集上的准确率。

关键点：

• swa_start： SWA的开始时间非常重要。通常，在训练的后期，模型已经接近收敛时启动SWA效果更好。
• 学习率： SWA通常需要一个比正常训练更小的学习率。
• 模型副本： SWA需要创建一个模型的副本 swa_model，用于存储平均后的权重。
• 权重更新： update_swa()函数负责更新SWA模型的权重，每次更新都会将当前模型的权重添加到SWA模型中，并计算平均值。

2.4 SWA的优点和缺点

优点：

• 简单易用： SWA的实现非常简单，只需要几行代码即可集成到现有的训练流程中。
• 计算效率高： SWA只需要训练一次模型，与模型集成相比，计算成本大大降低。
• 泛化能力强： SWA可以有效地提高模型的泛化能力，尤其是在损失函数曲面比较复杂的情况下。

缺点：

• 需要调整参数： SWA的性能对 swa_start 和学习率等参数比较敏感，需要进行调整。
• 可能需要更多内存： SWA需要保存多个模型的权重，可能需要更多的内存。

3. 指数移动平均（EMA）：一种在线平均策略

3.1 EMA的核心思想

EMA是一种在线平均策略，它在训练的每一步都对模型的权重进行指数加权平均。与SWA不同，EMA不需要保存多个模型的权重，而是维护一个移动平均的权重副本。

3.2 EMA的算法流程

1. 初始化： 创建一个模型权重的副本，初始化为模型的初始权重。

2. 权重更新： 在训练的每一步，使用以下公式更新EMA权重：

   EMA_weight = decay * EMA_weight + (1 - decay) * current_weight

   其中，decay 是一个介于0和1之间的衰减系数，用于控制EMA权重的更新速度。current_weight 是当前模型的权重。

3. 模型评估： 在训练结束后，使用EMA权重作为模型的最终权重进行评估。

3.3 EMA的Python实现

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

# 假设我们有一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

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

# 假设我们有一个自定义数据集
class DummyDataset(Dataset):
    def __init__(self, num_samples, input_size):
        self.num_samples = num_samples
        self.input_size = input_size
        self.data = torch.randn(num_samples, input_size)
        self.labels = torch.randint(0, 2, (num_samples,))  # 二分类问题

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

class EMA:
    def __init__(self, model, decay):
        self.model = model
        self.decay = decay
        self.shadow = {}
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def update(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.shadow[name].data = (
                    self.decay * self.shadow[name].data +
                    (1 - self.decay) * param.data
                )

    def apply_shadow(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.data, self.shadow[name].data = self.shadow[name].data, param.data

    def restore(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.data, self.shadow[name].data = self.shadow[name].data, param.data

def ema_implementation(model, train_loader, optimizer, num_epochs, decay, device):
    """
    EMA的实现函数。

    Args:
        model: PyTorch模型。
        train_loader: 训练数据加载器。
        optimizer: 优化器。
        num_epochs: 训练的epoch数。
        decay: EMA的衰减系数。
        device: 设备 (CPU 或 CUDA)。
    """
    ema = EMA(model, decay)

    for epoch in range(num_epochs):
        for i, (inputs, labels) in enumerate(train_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)

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

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

            # 更新EMA权重
            ema.update()

        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

    # 保存原始模型权重
    ema.apply_shadow()  # 使用EMA权重替换模型权重
    return model

# 超参数设置
input_size = 10
hidden_size = 20
output_size = 2
num_epochs = 20
batch_size = 32
learning_rate = 0.001
decay = 0.999 # EMA decay factor

# 创建模型、数据集和数据加载器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleModel(input_size, hidden_size, output_size).to(device)
dataset = DummyDataset(num_samples=1000, input_size=input_size)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 执行EMA
ema_model = ema_implementation(model, train_loader, optimizer, num_epochs, decay, device)

# 使用EMA模型进行推理 (示例)
def evaluate_model(model, data_loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

test_dataset = DummyDataset(num_samples=200, input_size=input_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 评估原始模型
# ema.restore() # restore original weights
accuracy_original = evaluate_model(model, test_loader, device)
print(f"原始模型准确率: {accuracy_original:.4f}")

# 评估EMA模型
# ema.apply_shadow()  # apply ema weights
accuracy_ema = evaluate_model(ema_model, test_loader, device)
print(f"EMA模型准确率: {accuracy_ema:.4f}")

代码解释：

• SimpleModel：一个简单的线性模型，用于演示EMA。
• DummyDataset：一个模拟数据集，用于训练模型。
• EMA：一个类，用于实现EMA的逻辑。它维护一个模型权重的副本 shadow，并使用指数加权平均更新该副本。
• ema_implementation：EMA的核心实现函数。它接收模型、数据加载器、优化器等参数，并在训练过程中应用EMA。
• evaluate_model：评估模型在测试集上的准确率。

关键点：

• decay： 衰减系数 decay 是EMA的关键参数。通常，decay 的值接近于1，例如0.999。
• shadow： shadow 是模型权重的副本，用于存储EMA权重。
• 更新： update()函数负责更新 shadow 中的权重，使用指数加权平均。
• apply_shadow 和 restore： apply_shadow()函数用于将EMA权重应用到模型中，restore()函数用于将原始权重恢复到模型中。

3.4 EMA的优点和缺点

优点：

• 内存效率高： EMA只需要维护一个模型权重的副本，内存效率比SWA更高。
• 在线平均： EMA是一种在线平均策略，可以在训练的每一步都对权重进行平均。
• 易于实现： EMA的实现相对简单，只需要几行代码即可集成到现有的训练流程中。

缺点：

• 需要调整参数： EMA的性能对 decay 参数比较敏感，需要进行调整。
• 可能导致训练不稳定： 如果 decay 值设置不当，EMA可能会导致训练不稳定。

4. SWA vs EMA：选择哪种方法？

SWA和EMA都是有效的模型参数平滑技术，但它们在适用场景和性能上有所不同。

特性	SWA	EMA
平均策略	离线平均，训练后期保存权重并平均	在线平均，每一步都更新权重
内存效率	较低，需要保存多个模型权重	较高，只需要保存一个模型权重的副本
参数调整	swa_start、学习率等	decay
适用场景	损失函数曲面比较复杂，需要更强的平均	内存资源有限，需要在线平均
实现复杂度	中等	简单

选择建议：

• 如果计算资源充足，并且希望获得更强的平均效果，可以尝试SWA。
• 如果内存资源有限，或者需要在训练过程中进行在线平均，可以选择EMA。
• 在实践中，可以尝试不同的方法，并根据实际情况选择最适合的模型参数平滑策略。

5. 总结：有效的参数平滑，提高模型泛化能力

我们深入探讨了两种主流的模型参数平滑技术：SWA和EMA。SWA通过在训练后期对模型权重进行离线平均，而EMA则通过在线指数加权平均来平滑模型参数。这两种方法都可以有效地提高模型的泛化能力和鲁棒性。选择哪种方法取决于具体的应用场景和计算资源。希望通过本文的讲解和代码示例，能够帮助大家更好地理解和应用模型参数平滑技术，提升深度学习模型的性能。

更多IT精英技术系列讲座，到智猿学院