📚 课程目录第1章嵌入式AI模型压缩与部署概述

1. 嵌入式AI模型压缩与部署概述

任务目标

实现一个基于PyTorch的权重剪枝流程：对预训练模型执行全局幅度剪枝（Global Magnitude Pruning），将指定比例的较小权重置零，并评估剪枝前后模型在测试集上的精度变化。本案例聚焦于剪枝流程本身，不涉及稀疏化加速或模型微调。

环境依赖

Python 3.8+
PyTorch 1.12+
torchvision

关键完整代码

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# 1. 加载预训练模型与测试数据
model = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
model.eval()

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=256, shuffle=False, num_workers=2)

# 2. 评估函数
def evaluate(model, loader):
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100.0 * correct / total

# 3. 全局幅度剪枝
def global_magnitude_prune(model, prune_ratio):
    # 收集所有卷积层和全连接层的权重
    all_weights = []
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:
            all_weights.append(param.data.abs().flatten())
    
    if not all_weights:
        return
    
    # 计算全局剪枝阈值
    all_flat = torch.cat(all_weights)
    k = int(all_flat.numel() * prune_ratio)
    threshold = torch.kthvalue(all_flat, k).values
    
    # 将低于阈值的权重置零
    with torch.no_grad():
        for name, param in model.named_parameters():
            if 'weight' in name and param.dim() >= 2:
                mask = param.data.abs() > threshold
                param.data.mul_(mask)

# 4. 执行剪枝与评估
prune_ratio = 0.5  # 剪枝50%的权重

acc_before = evaluate(model, testloader)
print(f"剪枝前准确率: {acc_before:.2f}%")

global_magnitude_prune(model, prune_ratio)

# 统计实际置零比例
total_weights = 0
zero_weights = 0
for name, param in model.named_parameters():
    if 'weight' in name and param.dim() >= 2:
        total_weights += param.numel()
        zero_weights += (param.data == 0).sum().item()
actual_ratio = zero_weights / total_weights * 100

acc_after = evaluate(model, testloader)
print(f"实际剪枝比例: {actual_ratio:.2f}%")
print(f"剪枝后准确率: {acc_after:.2f}%")
print(f"精度变化: {acc_after - acc_before:+.2f}个百分点")

运行方式

将代码保存为 pruning_demo.py；
直接运行 python pruning_demo.py（首次运行会自动下载CIFAR-10测试集与ResNet18预训练权重）；
程序依次输出剪枝前准确率、实际剪枝比例、剪枝后准确率与精度变化。

验证方法与判定标准

剪枝比例验证：程序输出的“实际剪枝比例”应接近设定的50%（由于阈值边界效应，允许±1%偏差）；
精度影响评估：剪枝后准确率应低于剪枝前，但通常仍保持在合理范围（例如CIFAR-10上ResNet18从约85%降至70%～80%区间，具体数值因随机性略有浮动）；
功能正确性：若实际剪枝比例显著偏离设定值（如为0%或100%），说明阈值计算或掩码应用存在错误。

常见问题

权重维度过滤：仅对维度≥2的权重（卷积核、全连接矩阵）剪枝，偏置和BatchNorm参数不参与，这是幅度剪枝的常见约定；
阈值边界：torch.kthvalue取第k小的值作为阈值，等于该值的权重会被保留（使用>而非>=），因此实际剪枝比例可能略低于设定值；
无微调：本案例仅演示剪枝与评估流程，不包含剪枝后微调（fine-tune）步骤，如需恢复精度需额外训练。