模型剪枝经典论文精读：NISP: Pruning Networks using Neuron Importance Score Propagation

一、论文基本信息

论文题目：NISP: Pruning Networks using Neuron Importance Score Propagation

作者：Ruichi Yu、Ang Li、Chun-Fu Chen、Jui-Hsin Lai、Vlad I. Morariu、Xintong Han、Mingfei Gao、Ching-Yung Lin、Larry S. Davis

发表信息：CVPR 2018

论文链接：

arXiv：NISP: Pruning Networks using Neuron Importance Score Propagation
CVF Open Access：CVPR 2018 论文页面
Adobe Research 页面：NISP: Pruning Networks using Neuron Importance Score Propagation

CVF 页面显示该论文收录于 Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, pp. 9194--9203 ，Adobe Research 页面也标注其为 CVPR 2018 Spotlight 工作。

这篇论文的核心问题是：很多剪枝方法只看单层或相邻两层，例如根据当前层权重大小剪 filter，或者根据下一层重构误差剪通道，但这些方法容易忽略深层网络中的误差传播。NISP 提出从 final response layer，FRL 出发，先衡量最终响应层中神经元的重要性，再把这些重要性分数反向传播到前面所有层，最终根据传播得到的重要性剪掉低分神经元或通道。

二、论文要解决的问题

在前面几篇论文中，我们已经看到几类典型剪枝思路。

Pruning Filters for Efficient ConvNets 直接用 filter 的 L1 范数判断重要性。

ThiNet 看下一层输出能否被重构。

Channel Pruning for Accelerating Very Deep Neural Networks 用 LASSO 选择通道，再用最小二乘重建权重。

Network Slimming 则通过 BN 的 (\gamma) 稀疏化来选择通道。

这些方法都有一个共同特点：它们大多是局部判断。要么只看当前层，要么只看当前层和下一层。NISP 认为，这种局部剪枝存在一个明显问题：

复制代码

早期层中某个神经元本身看起来不重要，
但它可能通过后续多层连接，
最终影响 final response layer 中非常重要的响应。

如果剪枝只看局部，就可能误删这种"间接重要"的神经元。因此，NISP 要解决的问题是：

如何从最终分类相关的响应出发，给整个网络中所有层的神经元分配统一的重要性分数？

论文明确指出，剪枝应该围绕一个统一目标：尽量减少 final response layer 中重要响应的重构误差，而不是逐层孤立地剪。

三、核心思想

NISP 的核心思想可以概括为一句话：

先在最终响应层 FRL 上计算神经元重要性，再把这个重要性沿着网络连接反向传播到前面各层，最后删除传播重要性最低的神经元或通道。

这里的 FRL 指的是分类层之前的倒数第二层，也就是直接输入分类器的特征层。论文认为，FRL 中的响应最接近最终分类任务，因此保留 FRL 中重要响应的表达能力，是剪枝后模型保持预测能力的关键。

整个流程可以写成：

复制代码

预训练 CNN
    ↓
在 FRL 上做 feature ranking
    ↓
得到 FRL 中每个神经元的重要性分数
    ↓
通过 NISP 将重要性分数向前传播
    ↓
得到每一层神经元 / 通道的重要性
    ↓
根据预设剪枝率删除低重要性神经元 / 通道
    ↓
fine-tuning 恢复精度

这和前面几篇论文最大的不同是：

复制代码

L1 Filter Pruning:
    当前层局部权重大小

ThiNet:
    下一层局部重构

Channel Pruning:
    当前层输出重构 + LASSO

Network Slimming:
    BN gamma 稀疏性

NISP:
    从最终响应层出发，向前传播重要性

所以 NISP 本质上是一种 global importance propagation 方法。

四、方法细节

4.1 什么是 Final Response Layer？

NISP 中最重要的概念是 Final Response Layer，FRL。FRL 是分类层之前的那一层特征响应。对于 CNN 分类网络来说，它通常是最后一个全连接层之前的特征层，或者全局池化后的特征层。

为什么选择 FRL？

因为 FRL 是分类器直接使用的特征表示。相比早期卷积层，FRL 更接近最终任务目标；相比最终 softmax 输出，FRL 又保留了更丰富的特征维度。因此，论文把 FRL 看作剪枝时的"全局目标层"。

论文中的直觉是：

复制代码

只要剪枝后还能较好保留 FRL 中的重要响应，
模型的最终分类能力就更容易保持。

因此，NISP 不直接问"某一层的某个通道在本层是否重要"，而是问：

复制代码

这个通道对 FRL 中的重要神经元有没有贡献？

这就是 NISP 与局部剪枝方法的根本区别。

4.2 第一步：在 FRL 上做特征排序

NISP 首先需要得到 FRL 中每个神经元的重要性分数。

论文把 FRL 中的神经元看作特征，然后使用 feature ranking 方法评估每个特征对分类任务的重要性。论文中采用的是 Infinite Feature Selection，Inf-FS，它可以综合考虑特征之间的方差和相关性，从而给 FRL 中每个神经元一个重要性分数。

这一步得到的是：

其中表示最终响应层，也就是第 (n) 层中所有神经元的重要性分数。

可以理解为：

复制代码

FRL 中哪些神经元最重要？
哪些神经元对分类最有贡献？

NISP 后续所有层的重要性都从这个出发向前传播。

4.3 第二步：把剪枝目标写成优化问题

NISP 的理论部分试图回答一个问题：

如果我想在第 (l) 层删掉一些神经元，应该保留哪些，才能尽量不影响 FRL 中的重要响应？

论文将这个问题写成一个二值优化问题。

NISP 的目标是：

复制代码

选择第 l 层要保留的神经元，
使剪枝前后 FRL 中重要响应的差异尽可能小。

论文通过一系列上界推导，把复杂目标转化成一个可求解的形式。最后得到的结论是：对于某一层，只要计算出该层每个神经元对 FRL 重要响应的传播贡献，然后保留分数最高的神经元即可。

这一步非常关键，因为它把"全局剪枝目标"变成了"重要性分数传播"。

4.4 第三步：Neuron Importance Score Propagation

NISP 最核心的公式是重要性传播公式。

论文定义第 (k) 层响应的重要性为：

其中：

：第 (k) 层神经元的重要性；
：FRL 的神经元重要性；
：第 (i) 层权重；
：权重逐元素取绝对值。

这个式子的意思是：

复制代码

某个早期神经元的重要性，
由它通过后续权重连接到 FRL 重要神经元的贡献决定。

论文进一步指出，这个重要性可以递归传播：

展开到单个神经元就是：

也就是说，第 (k) 层第 (j) 个神经元的重要性，等于它连接到后一层各个神经元的重要性加权和。后一层越重要、连接权重越大，那么当前神经元就越重要。论文明确指出，该传播形式也可以推广到 normalization、pooling 和 branch connections。

这就是 Neuron Importance Score Propagation 名字的来源。

4.5 第四步：卷积层如何剪？

对于全连接层，神经元就是一个输出维度，直接根据重要性排序删除即可。

对于卷积层，情况稍微不同。卷积层的一个 filter 会产生一个 feature map，也就是一个输出 channel。这个 channel 中包含多个空间位置上的神经元。

因此，论文把卷积层的一个 channel 作为剪枝单位：

复制代码

卷积层中一个 filter / output channel
    =
一个结构化剪枝单元

论文说明，对于卷积层，会整体剪掉一个 channel；该 channel 的重要性由这个 channel 内所有神经元的重要性求和得到。

也就是说：

然后根据 channel importance 排序，删除分数最低的通道。这样剪枝后得到的仍然是规则的窄网络，可以真正减少卷积计算量。

4.6 第五步：给定每层剪枝率，删除低分神经元

NISP 不是自动搜索每层剪枝率，而是假设每层剪枝率是预先给定的超参数。论文也明确说明，每层 pruning ratio 可以根据 FLOPs、内存和精度约束等实际需求来设置。

具体做法是：

复制代码

给定第 k 层需要保留 Nk 个神经元 / 通道
    ↓
根据 NISP 得到该层重要性分数 sk
    ↓
保留 top-Nk
    ↓
删除剩余低分神经元 / 通道

论文中还强调，重要性传播和层剪枝可以在一次反向遍历中联合完成；已经被剪掉的神经元，其重要性不会继续向更低层传播。

4.7 第六步：fine-tuning

剪枝后，NISP 会对剪枝后的网络进行 fine-tuning，以恢复由于删除神经元造成的性能损失。

这和大多数传统剪枝方法一致：

复制代码

预训练模型
    ↓
计算重要性
    ↓
剪枝
    ↓
fine-tuning

论文摘要和方法部分都明确说明，CNN 会删除低重要性神经元，然后通过 fine-tuning 恢复预测能力。

4.8 代码：

python 复制代码

"""
NISP VGG16-CIFAR10 Demo
============================================================

对应论文：
    NISP: Pruning Networks using Neuron Importance Score Propagation
    Ruichi Yu et al., CVPR 2018

这个文件是一个"理解性 + 可运行 demo"，不是官方代码逐行复现。

核心演示：
    1. 使用 CIFAR-10 数据集。
    2. 构建适配 CIFAR-10 的 VGG16-BN。
    3. 训练或加载一个 VGG16-CIFAR10 模型。
    4. 以最后分类器前的特征作为 FRL（Final Response Layer）。
    5. 用 classifier.weight 的绝对值和作为 FRL 通道重要性。
    6. 按照 NISP 思想：
           s_k = |W_{k+1}|^T s_{k+1}
       将重要性逐层从后往前传播到各个卷积层。
    7. 根据每层通道重要性剪掉低分通道。
    8. 物理构建一个更窄的 VGG16。
    9. 可选 fine-tuning 剪枝后的模型。

运行示例：

    # 快速 smoke test，训练 1 个 epoch 后剪枝
    python nisp_vgg16_cifar10_demo.py --epochs 1 --prune-ratio 0.3 --finetune-epochs 1

    # 只使用更少数据快速测试
    python nisp_vgg16_cifar10_demo.py --epochs 1 --subset 5000 --prune-ratio 0.3

    # 加载已有 checkpoint 后直接剪枝
    python nisp_vgg16_cifar10_demo.py --checkpoint vgg16_cifar10.pth --epochs 0 --prune-ratio 0.5

注意：
    - 第一次运行会下载 CIFAR-10。
    - 如果你只想理解 NISP 流程，epochs 可以设得很小。
    - 如果你想得到较好精度，需要正常训练更多 epoch。
"""

from __future__ import annotations

import argparse
from typing import Dict, List, Sequence, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
from torchvision import datasets, transforms


# ============================================================
# 1. VGG16-CIFAR10 模型
# ============================================================

VGG16_CFG: List[Union[int, str]] = [
    64, 64, "M",
    128, 128, "M",
    256, 256, 256, "M",
    512, 512, 512, "M",
    512, 512, 512, "M",
]


def make_vgg_layers(
    cfg: Sequence[Union[int, str]],
    in_channels: int = 3,
    batch_norm: bool = True,
) -> nn.Sequential:
    """
    根据 cfg 构建 VGG features。

    对 CIFAR-10 输入 32x32：
        经过 5 次 MaxPool2d 后，空间大小从 32 -> 16 -> 8 -> 4 -> 2 -> 1。
    """
    layers: List[nn.Module] = []

    for v in cfg:
        if v == "M":
            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
        else:
            out_channels = int(v)
            conv = nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=3,
                padding=1,
                bias=not batch_norm,
            )
            layers.append(conv)

            if batch_norm:
                layers.append(nn.BatchNorm2d(out_channels))

            layers.append(nn.ReLU(inplace=True))
            in_channels = out_channels

    return nn.Sequential(*layers)


class VGG16CIFAR(nn.Module):
    """
    适配 CIFAR-10 的 VGG16-BN。

    与 ImageNet VGG16 的区别：
        1. 输入是 32x32。
        2. 最后使用 AdaptiveAvgPool2d(1)。
        3. classifier 只有一个 Linear，而不是 4096-4096-1000。
    """

    def __init__(
        self,
        cfg: Sequence[Union[int, str]] = VGG16_CFG,
        num_classes: int = 10,
        batch_norm: bool = True,
    ) -> None:
        super().__init__()
        self.cfg = list(cfg)
        self.batch_norm = batch_norm

        self.features = make_vgg_layers(
            cfg=self.cfg,
            in_channels=3,
            batch_norm=batch_norm,
        )

        last_channels = [x for x in self.cfg if x != "M"][-1]
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Linear(int(last_channels), num_classes)

        self._init_weights()

    def _init_weights(self) -> None:
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        x = self.classifier(x)
        return x


def get_conv_bn_pairs(model: VGG16CIFAR) -> List[Tuple[nn.Conv2d, nn.BatchNorm2d]]:
    """
    取出 VGG features 中的 Conv-BN pair。
    """
    pairs: List[Tuple[nn.Conv2d, nn.BatchNorm2d]] = []
    modules = list(model.features)

    for i, m in enumerate(modules):
        if isinstance(m, nn.Conv2d):
            if i + 1 < len(modules) and isinstance(modules[i + 1], nn.BatchNorm2d):
                pairs.append((m, modules[i + 1]))
            else:
                raise RuntimeError("This demo expects Conv2d followed by BatchNorm2d.")

    return pairs


# ============================================================
# 2. CIFAR-10 数据
# ============================================================

def build_cifar10_loaders(
    data_root: str,
    batch_size: int,
    num_workers: int,
    subset: int = 0,
) -> Tuple[data.DataLoader, data.DataLoader]:
    """
    构建 CIFAR-10 train/test dataloader。

    subset > 0 时，只取训练集前 subset 个样本，方便快速调试。
    """
    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2470, 0.2435, 0.2616)

    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])

    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])

    train_set = datasets.CIFAR10(
        root=data_root,
        train=True,
        download=True,
        transform=train_transform,
    )

    test_set = datasets.CIFAR10(
        root=data_root,
        train=False,
        download=True,
        transform=test_transform,
    )

    if subset and subset > 0:
        indices = list(range(min(subset, len(train_set))))
        train_set = data.Subset(train_set, indices)

    train_loader = data.DataLoader(
        train_set,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
    )

    test_loader = data.DataLoader(
        test_set,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
    )

    return train_loader, test_loader


# ============================================================
# 3. 训练与评估
# ============================================================

def train_one_epoch(
    model: nn.Module,
    loader: data.DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    epoch: int,
    log_interval: int = 100,
) -> None:
    model.train()

    total_loss = 0.0
    total_correct = 0
    total_num = 0

    for step, (images, targets) in enumerate(loader, start=1):
        images = images.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        logits = model(images)
        loss = F.cross_entropy(logits, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        with torch.no_grad():
            pred = logits.argmax(dim=1)
            total_correct += (pred == targets).sum().item()
            total_loss += loss.item() * images.size(0)
            total_num += images.size(0)

        if step % log_interval == 0:
            print(
                f"epoch {epoch:03d} | step {step:04d}/{len(loader)} | "
                f"loss={total_loss / max(total_num, 1):.4f} | "
                f"acc={100.0 * total_correct / max(total_num, 1):.2f}%"
            )


@torch.no_grad()
def evaluate(
    model: nn.Module,
    loader: data.DataLoader,
    device: torch.device,
    title: str = "eval",
) -> Tuple[float, float]:
    model.eval()

    total_loss = 0.0
    total_correct = 0
    total_num = 0

    for images, targets in loader:
        images = images.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)

        logits = model(images)
        loss = F.cross_entropy(logits, targets)

        pred = logits.argmax(dim=1)
        total_correct += (pred == targets).sum().item()
        total_loss += loss.item() * images.size(0)
        total_num += images.size(0)

    avg_loss = total_loss / max(total_num, 1)
    acc = 100.0 * total_correct / max(total_num, 1)

    print(f"[{title}] loss={avg_loss:.4f}, acc={acc:.2f}%")
    return avg_loss, acc


# ============================================================
# 4. NISP：从 FRL 传播通道重要性
# ============================================================

@torch.no_grad()
def compute_nisp_channel_scores(model: VGG16CIFAR) -> List[torch.Tensor]:
    """
    计算每个卷积层输出通道的重要性分数。

    简化版 NISP：
        1. FRL 是 avgpool 后、classifier 前的特征。
        2. FRL 重要性用 classifier.weight 的绝对值求和：
               s_last = sum_class |W_fc[class, channel]|
        3. 对每个卷积层，从后往前传播：
               s_in = |W_conv|^T s_out
           对卷积核空间维求和后：
               W_channel[o, i] = sum_{kh,kw} |W[o, i, kh, kw]|
               s_in[i] = sum_o W_channel[o, i] * s_out[o]

    返回：
        scores[l] 是第 l 个 Conv2d 输出通道的重要性，shape [out_channels_l]。
    """
    pairs = get_conv_bn_pairs(model)
    convs = [conv for conv, _ in pairs]

    # FRL importance: classifier input channel importance.
    score = model.classifier.weight.detach().abs().sum(dim=0).cpu()

    scores_reversed: List[torch.Tensor] = []

    # 从最后一个 conv 往第一个 conv 传播。
    for conv in reversed(convs):
        # score 当前表示该 conv 输出通道的重要性。
        out_score = score.clone().float().cpu()
        scores_reversed.append(out_score)

        # 通过该 conv 的权重传播到该 conv 的输入通道。
        # conv.weight: [out_channels, in_channels, k, k]
        w = conv.weight.detach().abs().cpu()
        w_channel = w.sum(dim=(2, 3))  # [out_channels, in_channels]

        # s_in = |W|^T s_out
        score = w_channel.t().matmul(out_score)

    scores = list(reversed(scores_reversed))
    return scores


@torch.no_grad()
def make_keep_indices_from_scores(
    scores: List[torch.Tensor],
    prune_ratio: float,
    min_channels: int,
    skip_first_conv: bool = False,
) -> List[torch.Tensor]:
    """
    根据每层通道重要性分数，生成保留通道编号。

    prune_ratio:
        每层剪掉的比例。
        例如 0.5 表示每层保留约 50% 通道。

    min_channels:
        每层至少保留多少通道，防止剪空。

    skip_first_conv:
        是否不剪第一个卷积层输出。
        demo 中默认 False；如果希望更保守可以设 True。
    """
    keep_indices: List[torch.Tensor] = []

    for layer_idx, s in enumerate(scores):
        c = s.numel()

        if skip_first_conv and layer_idx == 0:
            keep = torch.arange(c).long()
        else:
            keep_count = int(round(c * (1.0 - prune_ratio)))
            keep_count = max(min_channels, keep_count)
            keep_count = min(c, keep_count)

            keep = torch.topk(s, k=keep_count, largest=True).indices
            keep = torch.sort(keep).values.long()

        keep_indices.append(keep)

    return keep_indices


def cfg_from_keep_indices(
    original_cfg: Sequence[Union[int, str]],
    keep_indices: Sequence[torch.Tensor],
) -> List[Union[int, str]]:
    """
    根据保留通道数生成新的 VGG cfg。
    """
    new_cfg: List[Union[int, str]] = []
    conv_idx = 0

    for v in original_cfg:
        if v == "M":
            new_cfg.append("M")
        else:
            new_cfg.append(int(keep_indices[conv_idx].numel()))
            conv_idx += 1

    return new_cfg


@torch.no_grad()
def copy_pruned_vgg_weights(
    old_model: VGG16CIFAR,
    new_model: VGG16CIFAR,
    keep_indices: Sequence[torch.Tensor],
) -> None:
    """
    把原始 VGG16 权重复制到剪枝后的 VGG16 中。

    对第 l 个卷积层：
        输出通道保留 keep_indices[l]
        输入通道保留 keep_indices[l-1]
        第一个卷积层的输入固定是 RGB [0, 1, 2]

    对 BN：
        同步复制 gamma、beta、running_mean、running_var。

    对 classifier：
        输入维度对应最后一个卷积层保留的通道。
    """
    old_pairs = get_conv_bn_pairs(old_model)
    new_pairs = get_conv_bn_pairs(new_model)

    in_keep = torch.arange(3).long()

    for layer_idx, ((old_conv, old_bn), (new_conv, new_bn)) in enumerate(zip(old_pairs, new_pairs)):
        out_keep = keep_indices[layer_idx].cpu().long()

        old_w = old_conv.weight.detach().cpu()
        new_w = old_w[out_keep][:, in_keep].clone()
        new_conv.weight.copy_(new_w)

        if old_conv.bias is not None and new_conv.bias is not None:
            new_conv.bias.copy_(old_conv.bias.detach().cpu()[out_keep])

        new_bn.weight.copy_(old_bn.weight.detach().cpu()[out_keep])
        new_bn.bias.copy_(old_bn.bias.detach().cpu()[out_keep])
        new_bn.running_mean.copy_(old_bn.running_mean.detach().cpu()[out_keep])
        new_bn.running_var.copy_(old_bn.running_var.detach().cpu()[out_keep])
        new_bn.num_batches_tracked.copy_(old_bn.num_batches_tracked.detach().cpu())

        in_keep = out_keep

    last_keep = keep_indices[-1].cpu().long()
    new_model.classifier.weight.copy_(old_model.classifier.weight.detach().cpu()[:, last_keep])
    new_model.classifier.bias.copy_(old_model.classifier.bias.detach().cpu())


@torch.no_grad()
def build_pruned_vgg16_by_nisp(
    old_model: VGG16CIFAR,
    prune_ratio: float = 0.5,
    min_channels: int = 8,
    skip_first_conv: bool = False,
    verbose: bool = True,
) -> Tuple[VGG16CIFAR, List[torch.Tensor], List[torch.Tensor]]:
    """
    根据 NISP 通道重要性构建剪枝后的 VGG16-CIFAR10。

    返回：
        pruned_model
        scores
        keep_indices
    """
    old_model_cpu = old_model.cpu().eval()

    scores = compute_nisp_channel_scores(old_model_cpu)
    keep_indices = make_keep_indices_from_scores(
        scores=scores,
        prune_ratio=prune_ratio,
        min_channels=min_channels,
        skip_first_conv=skip_first_conv,
    )

    new_cfg = cfg_from_keep_indices(old_model_cpu.cfg, keep_indices)
    pruned_model = VGG16CIFAR(
        cfg=new_cfg,
        num_classes=old_model_cpu.classifier.out_features,
        batch_norm=True,
    )

    copy_pruned_vgg_weights(old_model_cpu, pruned_model, keep_indices)

    if verbose:
        print("\n[NISP channel pruning plan]")
        conv_idx = 0
        for v in old_model_cpu.cfg:
            if v == "M":
                continue
            old_c = int(v)
            new_c = int(keep_indices[conv_idx].numel())
            s = scores[conv_idx]
            print(
                f"conv{conv_idx + 1:02d}: "
                f"{old_c:4d} -> {new_c:4d} | "
                f"score min={s.min().item():.4e}, "
                f"mean={s.mean().item():.4e}, "
                f"max={s.max().item():.4e}"
            )
            conv_idx += 1

    return pruned_model, scores, keep_indices


# ============================================================
# 5. 参数量与 MACs 统计
# ============================================================

def count_params(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


@torch.no_grad()
def profile_macs(
    model: nn.Module,
    input_size: Tuple[int, int, int, int],
    device: torch.device,
) -> Dict[str, float]:
    """
    简单统计 Conv2d 和 Linear 的 MACs。
    """
    model = model.to(device).eval()

    total_macs = 0
    conv_macs = 0
    linear_macs = 0
    handles = []

    def conv_hook(module: nn.Conv2d, inputs, output):
        nonlocal total_macs, conv_macs

        batch = output.shape[0]
        out_h, out_w = output.shape[-2:]
        kernel_ops = module.kernel_size[0] * module.kernel_size[1] * module.in_channels // module.groups
        macs = batch * module.out_channels * out_h * out_w * kernel_ops

        conv_macs += macs
        total_macs += macs

    def linear_hook(module: nn.Linear, inputs, output):
        nonlocal total_macs, linear_macs

        x = inputs[0]
        batch = x.shape[0] if x.dim() > 1 else 1
        macs = batch * module.in_features * module.out_features

        linear_macs += macs
        total_macs += macs

    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            handles.append(m.register_forward_hook(conv_hook))
        elif isinstance(m, nn.Linear):
            handles.append(m.register_forward_hook(linear_hook))

    dummy = torch.randn(*input_size, device=device)
    _ = model(dummy)

    for h in handles:
        h.remove()

    return {
        "params": float(count_params(model)),
        "macs": float(total_macs),
        "conv_macs": float(conv_macs),
        "linear_macs": float(linear_macs),
    }


def print_profile_compare(
    original: nn.Module,
    pruned: nn.Module,
    device: torch.device,
) -> None:
    s0 = profile_macs(original, (1, 3, 32, 32), device)
    s1 = profile_macs(pruned, (1, 3, 32, 32), device)

    def reduction(a: float, b: float) -> float:
        if a == 0:
            return 0.0
        return 100.0 * (a - b) / a

    print("\n[Model profile]")
    print(f"Params: {s0['params']:.0f} -> {s1['params']:.0f}, reduction={reduction(s0['params'], s1['params']):.2f}%")
    print(f"MACs:   {s0['macs']:.0f} -> {s1['macs']:.0f}, reduction={reduction(s0['macs'], s1['macs']):.2f}%")
    print(f"Conv:   {s0['conv_macs']:.0f} -> {s1['conv_macs']:.0f}, reduction={reduction(s0['conv_macs'], s1['conv_macs']):.2f}%")
    print(f"Linear: {s0['linear_macs']:.0f} -> {s1['linear_macs']:.0f}, reduction={reduction(s0['linear_macs'], s1['linear_macs']):.2f}%")


# ============================================================
# 6. Checkpoint
# ============================================================

def save_checkpoint(path: str, model: nn.Module, extra: Dict = None) -> None:
    ckpt = {
        "state_dict": model.state_dict(),
        "extra": extra or {},
    }
    torch.save(ckpt, path)
    print(f"Saved checkpoint to: {path}")


def load_checkpoint(path: str, model: nn.Module, device: torch.device) -> None:
    ckpt = torch.load(path, map_location=device)

    if isinstance(ckpt, dict) and "state_dict" in ckpt:
        state_dict = ckpt["state_dict"]
    else:
        state_dict = ckpt

    model.load_state_dict(state_dict)
    print(f"Loaded checkpoint from: {path}")


# ============================================================
# 7. Main
# ============================================================

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="NISP VGG16-CIFAR10 demo")

    parser.add_argument("--data-root", type=str, default="./data", help="CIFAR-10 root")
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--subset", type=int, default=0, help="Use first N training samples for quick demo")

    parser.add_argument("--epochs", type=int, default=1, help="Pre-pruning training epochs")
    parser.add_argument("--finetune-epochs", type=int, default=1, help="Post-pruning fine-tuning epochs")
    parser.add_argument("--lr", type=float, default=0.05)
    parser.add_argument("--finetune-lr", type=float, default=0.01)
    parser.add_argument("--weight-decay", type=float, default=5e-4)

    parser.add_argument("--prune-ratio", type=float, default=0.5, help="Per-layer channel pruning ratio")
    parser.add_argument("--min-channels", type=int, default=8)
    parser.add_argument("--skip-first-conv", action="store_true")

    parser.add_argument("--checkpoint", type=str, default="", help="Load original model checkpoint")
    parser.add_argument("--save-original", type=str, default="", help="Save original trained model")
    parser.add_argument("--save-pruned", type=str, default="", help="Save pruned model checkpoint")

    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")

    return parser.parse_args()


def main() -> None:
    args = parse_args()

    device = torch.device(args.device)
    print(f"Using device: {device}")

    train_loader, test_loader = build_cifar10_loaders(
        data_root=args.data_root,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        subset=args.subset,
    )

    model = VGG16CIFAR(num_classes=10, batch_norm=True).to(device)

    if args.checkpoint:
        load_checkpoint(args.checkpoint, model, device)

    if args.epochs > 0:
        optimizer = torch.optim.SGD(
            model.parameters(),
            lr=args.lr,
            momentum=0.9,
            weight_decay=args.weight_decay,
        )
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(args.epochs, 1))

        for epoch in range(1, args.epochs + 1):
            train_one_epoch(model, train_loader, optimizer, device, epoch)
            evaluate(model, test_loader, device, title=f"original after epoch {epoch}")
            scheduler.step()

    evaluate(model, test_loader, device, title="original before pruning")

    if args.save_original:
        save_checkpoint(
            args.save_original,
            model,
            extra={"cfg": model.cfg, "type": "VGG16CIFAR"},
        )

    # NISP pruning
    print("\nRunning NISP importance propagation and channel pruning...")
    pruned_model, scores, keep_indices = build_pruned_vgg16_by_nisp(
        old_model=model,
        prune_ratio=args.prune_ratio,
        min_channels=args.min_channels,
        skip_first_conv=args.skip_first_conv,
        verbose=True,
    )
    pruned_model = pruned_model.to(device)

    print_profile_compare(model, pruned_model, device)

    evaluate(pruned_model, test_loader, device, title="pruned before fine-tuning")

    if args.finetune_epochs > 0:
        optimizer = torch.optim.SGD(
            pruned_model.parameters(),
            lr=args.finetune_lr,
            momentum=0.9,
            weight_decay=args.weight_decay,
        )
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(args.finetune_epochs, 1))

        for epoch in range(1, args.finetune_epochs + 1):
            train_one_epoch(pruned_model, train_loader, optimizer, device, epoch)
            evaluate(pruned_model, test_loader, device, title=f"pruned after finetune epoch {epoch}")
            scheduler.step()

    if args.save_pruned:
        save_checkpoint(
            args.save_pruned,
            pruned_model,
            extra={
                "cfg": pruned_model.cfg,
                "type": "NISP-pruned VGG16CIFAR",
                "prune_ratio": args.prune_ratio,
                "min_channels": args.min_channels,
                "keep_indices": [k.tolist() for k in keep_indices],
            },
        )

    print("\nDone.")
    print("Key idea:")
    print("  FRL importance comes from classifier weights.")
    print("  NISP propagates importance backward through absolute conv weights.")
    print("  Low-score channels are pruned and a narrower VGG is physically built.")


if __name__ == "__main__":
    main()

五、关键公式

5.1 网络表示

论文将网络表示为多个层函数的复合：

其中第 (l) 层可写为：

5.2 FRL 重要性

首先在最终响应层得到重要性：

其中 (n) 表示 FRL。

5.3 全路径重要性传播

第 (k) 层的重要性可以写为：

这表示从 FRL 的重要性出发，沿权重连接向前传播。

5.4 递归传播公式

NISP 实际采用递归形式：

单个神经元形式为：

这个公式是整篇论文最核心的公式。它说明：一个低层神经元的重要性，来自它对后续重要神经元的连接贡献。

六、实验设置

论文在三个分类数据集上实验：

复制代码

MNIST
CIFAR-10
ImageNet

使用的网络包括：

复制代码

LeNet
Cifar-net
AlexNet
GoogLeNet
ResNet

论文中的实验和时间 benchmark 都基于 Caffe 完成；作者还使用 PCA accumulated energy analysis 来辅助选择剪枝率。

实验主要比较了几类 baseline：

复制代码

1. Random pruning
2. Train from scratch
3. Magnitude-based pruning
4. Layer-by-layer pruning
5. Existing pruning/compression methods

其中 random pruning 和 train-from-scratch 用来验证 NISP 的重要性排序是否有效；magnitude-based pruning 用来比较 FRL 重要性计算方式；layer-by-layer pruning 用来说明 NISP 这种全局传播比逐层独立剪枝更合理。

七、实验结果解读

7.1 MNIST / LeNet：50% 剪枝几乎不掉点

在 LeNet on MNIST 实验中，论文剪掉全连接层一半神经元，并剪掉两个卷积层一半 filters。结果显示，NISP 在 50% pruning ratio 下，相比剪枝前网络只降低 0.02% top-1 accuracy。

这个结果说明，在简单网络上，NISP 可以很准确地区分重要神经元和冗余神经元。

更重要的是，论文发现 NISP 在 fine-tuning 开始前的初始精度损失就比随机剪枝和从头训练更小，说明它删除的确实是对最终响应影响较小的部分。

7.2 CIFAR-10 / Cifar-net：50% 剪枝后损失小于 1%

在 Cifar-net on CIFAR-10 上，论文同样使用 50% pruning ratio。实验显示，NISP 在 fine-tuning 后的 top-1 accuracy loss 小于 1%。

这说明 NISP 不只是适用于 MNIST 这种简单数据集，在更复杂的自然图像分类任务上也能有效。

论文还指出，NISP 相比随机剪枝和从头训练，有更低的初始精度损失、更快的收敛速度，以及更低的最终精度损失。

7.3 ImageNet / AlexNet：同时剪 conv 和 FC

在 AlexNet on ImageNet 上，论文测试了两种剪枝方式：

复制代码

NISP-C:
    只剪 convolution layers

NISP-CF:
    同时剪 convolution layers 和 fully connected layers

论文对所有卷积层和全连接层采用 50% pruning ratio，并 fine-tune 90 epochs，报告 top-5 accuracy loss。结果显示，在只剪卷积层和同时剪卷积层/全连接层两种情况下，NISP 都保留了前面小数据集实验中观察到的优势。

这个实验很重要，因为 AlexNet 的全连接层参数量非常大。NISP 同时支持卷积层和全连接层剪枝，所以能实现真正的全网络压缩，而不仅仅是卷积加速。

7.4 ImageNet / GoogLeNet：处理 Inception 分支结构

GoogLeNet 中有 Inception 模块，包含多分支结构。论文在 GoogLeNet 上设计了多种剪枝策略，例如：

复制代码

Half:
    所有 convolution layers 都剪一半

noReduce:
    不剪 Inception 模块中的 reduction layers

no1x1:
    不剪 Inception 模块中的 1×1 convolution layers

论文在 GoogLeNet 上 fine-tune 60 epochs，并报告 top-5 accuracy loss。实验结果显示，NISP 在这种多分支网络上也能工作。

这说明 NISP 的传播思想不仅适用于简单串行网络，也可以扩展到带 branch connections 的结构。

7.5 与 magnitude-based pruning 的比较

论文专门比较了两种 FRL 重要性初始化方式：

复制代码

NISP-FS:
    使用 feature selection 方法计算 FRL 重要性

NISP-Mag:
    使用权重幅值计算 FRL 重要性

实验发现，NISP-FS 的精度损失明显小于 NISP-Mag，但 NISP-Mag 仍然优于 random pruning 和 train-from-scratch baseline。

这个结果说明两点。

第一，FRL 上的 feature ranking 很重要。

第二，即使用较简单的权重幅值作为 FRL 重要性，只要经过 NISP 向前传播，也比完全局部剪枝更合理。

7.6 与 Layer-by-Layer 剪枝的比较

Layer-by-Layer baseline 是指：每一层单独做 feature ranking，然后独立剪枝。

论文发现，虽然 Layer-by-Layer 比随机剪枝好，但它明显差于 NISP。原因是逐层独立剪枝无法考虑误差在深层网络中的传播。

论文还提出了 WARE，即 Weighted Average Reconstruction Error，用于衡量剪枝前后 FRL 中重要响应的变化。实验显示，当网络更深、剪枝比例更大时，Layer-by-Layer 的 WARE 会显著上升，而 NISP 更稳定。

这正好验证了 NISP 的核心动机：

复制代码

深层网络中，局部剪枝会导致误差传播；
必须从最终响应层出发做全局重要性传播。

7.7 与已有方法的压缩 benchmark

论文 Table 1 给出了多个模型的压缩结果。例如：

复制代码

AlexNet on ImageNet:
    NISP-A: accuracy loss 1.43%, FLOPs reduction 67.85%, params reduction 33.77%
    NISP-B: accuracy loss 0.97%, FLOPs reduction 62.69%, params reduction 1.96%
    NISP-C: accuracy loss 0.54%, FLOPs reduction 53.70%, params reduction 2.91%
    NISP-D: accuracy loss 0.00%, FLOPs reduction 40.12%, params reduction 47.09%

GoogLeNet on ImageNet:
    NISP: accuracy loss 0.21%, FLOPs reduction 58.34%, params reduction 33.76%

ResNet on CIFAR-10:
    NISP-56: accuracy loss 0.03%, FLOPs reduction 43.61%, params reduction 42.60%
    NISP-110: accuracy loss 0.18%, FLOPs reduction 43.78%, params reduction 43.25%

ResNet on ImageNet:
    NISP-34-B: accuracy loss 0.92%, FLOPs reduction 43.76%, params reduction 43.68%
    NISP-50-B: accuracy loss 0.89%, FLOPs reduction 44.01%, params reduction 43.82%

这些结果说明，NISP 不只是理论上提出了一个传播分数，也确实能在 AlexNet、GoogLeNet、ResNet 等网络上实现明显 FLOPs 和参数量下降，同时保持较小精度损失。

八、方法优点

8.1 从局部剪枝走向全局重要性传播

NISP 最大的贡献是把剪枝重要性从"本层局部指标"提升到了"面向最终响应层的全局指标"。

它不再只问：

复制代码

这个 filter 权重大不大？
这个通道能不能重构下一层？

而是问：

复制代码

这个神经元是否连接到最终响应层中的重要神经元？

这使得它比单层或相邻两层剪枝更适合深层网络。

8.2 理论动机比较清楚

NISP 不是直接提出一个经验公式，而是从"最小化 FRL 重要响应重构误差"出发，推导出重要性传播公式。

虽然最终得到的是一个近似解，但它给出了比较清晰的解释：

复制代码

低层神经元的重要性 = 后续重要神经元的重要性加权和

这比单纯用权重范数或 BN gamma 更具有全局解释性。

8.3 支持多种层类型

论文说明，NISP 可以应用于全连接层、卷积层，也可以扩展到 normalization、pooling 和 branch connections。对于卷积层，它剪掉整个 channel；对于全连接层，它剪掉单个 neuron。

因此，NISP 可以实现全网络压缩，而不是只处理某一类层。

8.4 可部署性好

NISP 输出的是更小的同类型网络，而不是非结构化稀疏矩阵。论文也指出，给定预训练模型，NISP 输出的是一个 smaller network of the same type，可以部署在原模型对应的硬件设备上。

这意味着它不依赖特殊稀疏计算库，更接近真实推理加速。

九、方法局限

9.1 每层剪枝率仍然需要预设

NISP 的重要性分数可以告诉我们"这一层哪些神经元更重要"，但它并没有自动决定每层应该剪多少。

论文中每层 pruning ratio 是预定义超参数。虽然作者说可以根据 FLOPs、内存和精度约束设置，但这仍然需要人工调参。

也就是说，NISP 解决的是：

复制代码

给定每层剪枝率，剪哪些神经元？

而不是：

复制代码

每层应该剪多少？

这一点和后来的自动剪枝、硬件感知剪枝相比仍有差距。

9.2 依赖 FRL 的特征排序质量

NISP 的所有传播都从 FRL 重要性 (s_n) 开始。如果 FRL 上的 feature ranking 不准确，后面传播得到的重要性也会受到影响。

论文采用 Inf-FS，但这也引入了额外算法和超参数，例如 loading coefficient (\alpha)。

因此，NISP 的效果不仅取决于传播公式，也取决于 FRL 重要性初始化是否可靠。

9.3 传播公式仍然是近似

NISP 使用绝对权重传播重要性：

这个公式简单高效，但它本质上仍然是近似。它没有完整考虑激活分布、非线性饱和、数据依赖梯度、BN 统计变化等因素。

因此，它比纯权重范数更全局，但仍不是严格等价于"删除该神经元造成的真实任务损失"。

9.4 对现代 Transformer 不直接适用

NISP 是为 CNN 中的 neuron / filter / channel pruning 设计的。对于 ViT、LLM、VLM 等模型，剪枝对象可能变成：

复制代码

attention heads
MLP hidden neurons
tokens
layers
KV cache
vision tokens

NISP 的"从最终响应层反向传播重要性"的思想可以借鉴，但原始公式不能直接照搬。

十、后续影响

NISP 在剪枝研究中的意义主要体现在三个方面。

第一，它强调剪枝不能只看局部层。深层网络中，早期层的某些通道可能通过后续连接影响最终任务，因此需要全局重要性度量。

第二，它把 feature selection 和 network pruning 联系起来。FRL 重要性来自 feature ranking，然后再向前传播到全网络。这让剪枝从单纯参数压缩变成了"任务相关特征保留"。

第三，它为后续重要性传播、全局剪枝、图依赖剪枝提供了思路。虽然 NISP 本身仍然需要预设剪枝率，但它已经开始把剪枝看作一个全网络联合问题，而不是逐层独立排序问题。

从专栏脉络上看，这篇论文可以放在这里：

复制代码

Pruning Filters for Efficient ConvNets
    ↓
ThiNet
    ↓
Channel Pruning
    ↓
Network Slimming
    ↓
NISP
    ↓
Rethinking the Value of Network Pruning
    ↓
Taylor Pruning / HRank / EagleEye / AutoSlim

如果说 ThiNet 和 Channel Pruning 主要回答的是：

复制代码

如何根据局部重构选择通道？

那么 NISP 回答的是：

复制代码

如何根据最终响应的重要性，给整个网络的神经元分配全局重要性？

十一、一句话总结

《NISP: Pruning Networks using Neuron Importance Score Propagation》从最终响应层 FRL 出发，通过权重连接将重要性分数向前传播到整个网络，再根据传播得到的神经元/通道重要性进行结构化剪枝，是早期"全局重要性传播剪枝"路线的代表工作。