大模型边端化的结构化剪枝实践：基于 L1 范数的通道剪枝与微调恢复完整流程

一、2GB 模型在边缘显存上的"硬着陆"：权重复制比参数更致命

大模型（如 LLaMA 3B、MobileLLaMA 1.5B）在 PC 端推理尚可应付，但迁移到 RK3588（8GB）或 Jetson Orin Nano（4GB）上时，模型权重本身加上推理时的 KV Cache、中间激活值，整体内存占用轻松突破 6GB。仅凭 int8 量化能将权重从 FP16 的 3GB 降至 1.5GB，但激活值和 cache 占用依然不变。

结构化剪枝与非结构化剪枝的核心差异在于：非结构化剪枝产生稀疏矩阵，需要专门的稀疏计算库支持，在通用 GPU/NPU 上的加速效果有限；而结构化剪枝以整通道（Channel）或整层为粒度进行裁剪，剪枝后的模型仍然是稠密的，可以直接在现有推理框架中以更小的形状运行，无需稀疏计算支持。本文以基于 L1 范数的通道重要性评估方法为核心，展示完整剪枝—微调闭环。

二、通道重要性评估与剪枝策略：L1 范数排序的数学依据

对于卷积层的输出通道，L1 范数衡量的是该通道滤波器权重的总幅度。范数越小的通道，对输出的贡献越低——这是剪枝的启发式基础。

flowchart TD
    A["原始模型<br/>Conv: 64→128 通道"] --> B["逐通道计算 L1 范数<br/>||W_i||₁ = Σ|w|"]
    B --> C{"范数排序<br/>升序排列"}
    C --> D["选取底部 K% 通道<br/>（剪枝率 30% → 保留 70%）"]
    D --> E["裁减权重矩阵<br/>W: [128,64,K,K] → [90,64,K,K]"]
    E --> F["裁减下一层对应通道<br/>W_next: [256,128,K,K] → [256,90,K,K]"]
    F --> G["结构化剪枝后模型<br/>（仍为稠密矩阵）"]
    G --> H["微调恢复精度<br/>10-20 轮训练"]
    H --> I["部署验证<br/>精度下降 < 2%"]

    C -->|"通道间范数方差分析"| J["自动确定剪枝率<br/>（基于方差拐点）"]
    J --> D

结构化剪枝的约束在于"通道对应"关系：剪掉第 L 层的某个输出通道，必须同时剪掉第 L+1 层对应的输入通道。这要求建立层间依赖图以确保结构兼容性。

三、PyTorch 实现：自动通道依赖图构建与剪枝执行

以下代码以 ResNet 的 BasicBlock 为例，展示通道依赖关系的自动追踪和剪枝执行：

# channel_pruning.py —— 基于 L1 范数的结构化通道剪枝

import torch
import torch.nn as nn
import numpy as np
from collections import OrderedDict
from typing import Dict, List, Tuple

# ——— 依赖图构建：在相邻层之间建立通道对应关系 ———
class DependencyGraph:
    """
    记录每层的输入/输出通道依赖关系。
    关键规则：
    - Conv → BN → ReLU：剪枝 Conv 输出 = 剪枝 BN 通道
    - Conv → Conv：剪枝前一层输出 = 剪枝后一层输入
    - Add/Concat (残差连接)：分支的输出通道数必须一致
    """

    def __init__(self, model: nn.Module):
        self.model = model
        # 记录每层名称到其剪枝掩码的映射
        self.layer_masks: Dict[str, torch.Tensor] = {}

    def build_graph(self) -> Dict[str, List[str]]:
        """
        返回依赖关系字典：
        { 'layer1.conv1': ['layer1.bn1'], 'layer1.bn1': ['layer1.conv2'], ... }
        """
        dep_graph = OrderedDict()
        prev_layer = None

        for name, module in self.model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                if prev_layer is not None:
                    dep_graph.setdefault(prev_layer, []).append(name)
                prev_layer = name
            elif isinstance(module, nn.BatchNorm2d):
                if prev_layer is not None:
                    dep_graph.setdefault(prev_layer, []).append(name)
                prev_layer = name

        return dep_graph

# ——— L1 范数重要性评估 ———
def l1_channel_importance(weight: torch.Tensor) -> torch.Tensor:
    """
    计算卷积/全连接层每个输出通道的 L1 范数。
    weight shape: [out_channels, in_channels, kernel_h, kernel_w] (Conv)
                  [out_features, in_features]                     (Linear)

    返回: shape [out_channels] 或 [out_features] 的重要性得分
    """
    # 沿着除输出通道外的所有维度求和
    dims = tuple(range(1, weight.dim()))  # (1,2,3) for Conv, (1,) for Linear
    return weight.abs().sum(dim=dims)

# ——— 结构化剪枝执行函数 ———
def structured_channel_prune(
    model: nn.Module,
    dep_graph: DependencyGraph,
    prune_ratio: float = 0.3
) -> nn.Module:
    """
    对模型中所有 Conv2d 和 Linear 层执行结构化通道剪枝。

    @param model:       原始 PyTorch 模型
    @param dep_graph:   层间依赖图
    @param prune_ratio: 剪枝比例（0.3 = 保留 70% 通道）
    @return: 剪枝后的模型
    """
    model_copy = copy.deepcopy(model)  # 避免修改原始模型
    masks: Dict[str, torch.Tensor] = {}

    for name, module in model_copy.named_modules():
        if not isinstance(module, (nn.Conv2d, nn.Linear)):
            continue

        # —— 步骤1：计算重要性并生成掩码 ——
        importance = l1_channel_importance(module.weight.data)
        num_channels = importance.size(0)
        num_keep = int(num_channels * (1.0 - prune_ratio))

        # 最小保留 1 个通道，避免层消失
        num_keep = max(num_keep, 1)

        # 获取重要性最低的通道索引（按得分升序取头部）
        _, sorted_indices = torch.sort(importance)
        keep_indices = sorted_indices[-num_keep:]  # 保留得分最高的 K 个

        # 生成布尔掩码 [out_channels]
        mask = torch.zeros(num_channels, dtype=torch.bool)
        mask[keep_indices] = True
        masks[name] = mask

    # —— 步骤2：应用掩码，裁剪权重矩阵 ——
    # 需要处理依赖关系：层A 的 mask 同时影响层A 的输出和层B 的输入
    for name, module in model_copy.named_modules():
        if isinstance(module, nn.Conv2d):
            out_mask = masks.get(name)
            if out_mask is None:
                continue

            # 裁剪输出通道
            module.weight.data = module.weight.data[out_mask, :, :, :]
            if module.bias is not None:
                module.bias.data = module.bias.data[out_mask]
            module.out_channels = int(out_mask.sum().item())

        elif isinstance(module, nn.BatchNorm2d):
            out_mask = masks.get(name)
            if out_mask is None:
                continue
            module.weight.data = module.weight.data[out_mask]
            module.bias.data = module.bias.data[out_mask]
            module.running_mean.data = module.running_mean.data[out_mask]
            module.running_var.data = module.running_var.data[out_mask]
            module.num_features = int(out_mask.sum().item())

    # —— 步骤3：裁剪下一层的输入通道（依赖传播） ——
    # 这步遍历依赖图，确保下一层输入通道与上一层输出对齐
    for layer_name, children in dep_graph.items():
        out_mask = masks.get(layer_name)
        if out_mask is None:
            continue

        for child_name in children:
            child_module = dict(model_copy.named_modules()).get(child_name)
            if child_module is None:
                continue

            if isinstance(child_module, nn.Conv2d):
                # Conv 的输入通道需要与上一层输出通道对齐
                if child_module.weight.data.size(1) != out_mask.size(0):
                    # 形状不一致（如 first conv 从 3 通道开始），跳过
                    continue
                child_module.weight.data = child_module.weight.data[:, out_mask, :, :]
                child_module.in_channels = int(out_mask.sum().item())

    return model_copy

# ——— 微调恢复流程 ———
def finetune_recover(
    pruned_model: nn.Module,
    train_loader: torch.utils.data.DataLoader,
    epochs: int = 15,
    lr: float = 1e-4
) -> nn.Module:
    """
    对剪枝后的模型进行少量轮次微调，恢复精度。
    建议：学习率设为原始训练的 1/10，避免破坏已保留的权重结构。
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pruned_model = pruned_model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    # 使用小学习率 + 权重衰减防止过拟合
    optimizer = torch.optim.AdamW(
        pruned_model.parameters(),
        lr=lr,
        weight_decay=1e-4
    )

    pruned_model.train()
    for epoch in range(epochs):
        total_loss = 0.0
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)

            optimizer.zero_grad()
            outputs = pruned_model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(train_loader)
        print(f"[Epoch {epoch+1}/{epochs}] loss={avg_loss:.4f}")

    return pruned_model

四、边界分析：剪枝率的精度-加速非线性关系与跨层方差

L1 范数剪枝的精度变化呈非线性特征。在 ResNet-50 的 ImageNet 测试中：

跨层方差问题：不同层对剪枝的敏感度差异可达 10 倍。浅层卷积（如 model.conv1）提取边缘和纹理特征，通道间信息冗余低；深层特征图通道间相关性高。因此固定全模型统一剪枝率的做法是次优的——理想方案是每层独立评估敏感度后分配差异化的剪枝率。

稀疏矩阵加速的误区：非结构化剪枝虽然参数减少比例更高，但在 GPU/NPU 上需要 CSR/CSC 格式的稀疏矩阵乘。以 NCNN 为例，稀疏卷积的实现仅对手机 GPU（Adreno/Mali）做了优化，在 Mali-G52 上稀疏加速约 1.3x，远低于理论值。因此边缘设备上优先选择结构化剪枝。

五、总结

结构化通道剪枝通过 L1 范数评估通道重要性，保留稠密矩阵特性，使得剪枝后的模型可以直接加载到 NCNN、MNN 等边缘推理框架中，无需依赖稀疏计算后端。在 30% 以下剪枝率范围内，配合 15 轮微调即可将精度损失控制在 0.5% 以内。关键实践点有三：构建准确的通道依赖图以确保结构兼容性；根据层敏感度分配差异化的剪枝率而非一刀切；微调时使用原始训练 1/10 量级的学习率以保护已保留的权重结构不受破坏。