3. 基于BN层缩放因子的通道稀疏化训练

任务目标

在PyTorch中实现基于BN层缩放因子的通道稀疏化训练：对卷积层后接的BatchNorm2d的gamma参数施加L1正则化，使不重要的通道gamma趋近于零，训练完成后可直接按gamma值裁剪通道，为后续模型压缩做准备。

核心原理

BN层每个通道对应一个可学习的缩放因子gamma（weight），其大小反映该通道对输出的贡献程度。在训练损失中加入 λ·|γ| 稀疏正则项，梯度下降会推动冗余通道的gamma逼近0，从而实现通道级稀疏。

环境依赖

  • Python 3.8+，PyTorch 1.10+，torchvision
  • 无需额外库，仅需标准PyTorch接口

关键完整代码

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

# ---------- 1. 稀疏训练核心：对BN的gamma施加L1正则 ----------
def sparse_loss(model, lambda_l1=1e-4):
    """计算所有BatchNorm2d层gamma的L1范数之和"""
    l1 = 0.0
    for module in model.modules():
        if isinstance(module, nn.BatchNorm2d):
            l1 += torch.abs(module.weight).sum()  # weight即gamma
    return lambda_l1 * l1

# ---------- 2. 简单CNN模型（含BN层） ----------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        self.classifier = nn.Linear(32, 10)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# ---------- 3. 训练循环（含稀疏正则） ----------
def train_one_epoch(model, loader, optimizer, criterion, lambda_l1):
    model.train()
    total_loss = 0.0
    for inputs, targets in loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        ce_loss = criterion(outputs, targets)
        reg_loss = sparse_loss(model, lambda_l1)  # 稀疏正则项
        loss = ce_loss + reg_loss
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

# ---------- 4. 验证gamma稀疏度 ----------
def check_sparsity(model, threshold=1e-4):
    """统计gamma值小于阈值的通道比例"""
    total, zeros = 0, 0
    for module in model.modules():
        if isinstance(module, nn.BatchNorm2d):
            g = module.weight.detach().abs()
            total += g.numel()
            zeros += (g < threshold).sum().item()
    return zeros / total if total > 0 else 0.0

# ---------- 5. 最小可运行示例 ----------
if __name__ == "__main__":
    # 构造随机数据模拟训练（实际使用时替换为真实数据集）
    torch.manual_seed(42)
    model = SimpleCNN()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)
    criterion = nn.CrossEntropyLoss()
    dummy_loader = [(torch.randn(8, 3, 32, 32), torch.randint(0, 10, (8,))) for _ in range(5)]

    lambda_l1 = 1e-4  # 稀疏强度超参数
    print(f"初始稀疏度: {check_sparsity(model):.4f}")
    for epoch in range(3):
        loss = train_one_epoch(model, dummy_loader, optimizer, criterion, lambda_l1)
        sparsity = check_sparsity(model)
        print(f"Epoch {epoch+1}: loss={loss:.4f}, 稀疏度={sparsity:.4f}")
    print("训练完成")

运行方式

  1. 将上述代码保存为 sparse_training.py
  2. 直接运行：python sparse_training.py
  3. 若需真实数据，将 dummy_loader 替换为CIFAR-10等数据集的DataLoader即可

验证方法与判定标准

  • 稀疏度指标：运行代码观察输出的“稀疏度”数值。该值表示gamma值小于1e-4的通道占比。训练前接近0，训练后应显著上升（例如从0%升至30%~80%，具体取决于λ和训练轮数）。
  • gamma分布可视化：训练后打印 model.features[1].weight.data（第一个BN层的gamma），可观察到大量接近0的值。
  • 通过条件：稀疏度从初始的0%提升至20%以上，且分类损失未发散（仍在正常下降或收敛），即证明L1正则成功实现了通道稀疏化。

关键参数说明

  • λ（lambda_l1）：稀疏强度。λ越大，gamma越稀疏，但可能损害精度。建议范围1e-5~1e-3，需根据任务调整。
  • weight_decay：SGD自带的L2正则不影响gamma稀疏，可保留用于权重衰减。
  • 裁剪时机：本案例只完成稀疏化训练，不包含裁剪步骤。训练完成后，可依据gamma值（如值小于阈值）移除对应通道及后续层输入维度。

常见问题

  • 稀疏度不上升：检查λ是否过小，或学习率是否过大导致gamma震荡。可增大λ或降低学习率。
  • 精度骤降：λ过大导致过多通道被稀疏化。需调小λ，或训练后期逐步增大λ（渐进式稀疏）。
  • BN层gamma初始全为1：这是PyTorch默认初始化，L1正则会在训练中逐步将不重要的gamma推向0。
← 上一章 📖 返回目录 下一章 →