YOLOv12模型剪枝实战：使用网络瘦身工具提升推理速度

最近在部署一个目标检测项目到边缘设备上，模型用的是最新的YOLOv12，效果确实不错，但推理速度总是不太理想。设备资源有限，跑起来有点吃力。相信不少朋友也遇到过类似问题：模型精度上去了，但速度下来了，在资源受限的设备上部署就成了大难题。

这时候，模型剪枝就派上用场了。简单来说，剪枝就是给模型“瘦身”，把那些对最终结果影响不大的部分去掉，让模型变得更轻巧，跑得更快。今天我就结合自己的实际经验，详细聊聊怎么用常见的剪枝工具，给YOLOv12模型做一次有效的“瘦身”，在保证精度损失不大的前提下，大幅提升推理速度。

1. 为什么需要给YOLOv12模型剪枝？

YOLOv12作为目标检测领域的佼佼者，精度和性能都很出色。但它的网络结构相对复杂，参数量和计算量都不小。当我们想把这样的模型部署到手机、嵌入式设备或者一些算力有限的边缘计算盒子时，就会遇到瓶颈。

直接部署原版模型，推理一帧图片可能要几百毫秒，甚至更久，这显然无法满足实时性要求高的场景，比如视频监控、自动驾驶感知或者移动端应用。模型剪枝的核心思路，就是找出网络中那些“冗余”的部分。你可以把神经网络想象成一棵枝繁叶茂的大树，有些枝叶对开花结果（也就是模型输出）贡献很小，甚至没有贡献。剪枝就是把这些“无效枝叶”修剪掉，让养分（计算资源）更集中地输送给重要的部分。

经过剪枝的模型，参数量（模型文件大小）和计算量（FLOPs）都会显著下降。这意味着模型加载更快、占用内存更少、推理速度更快，非常适合在资源受限的环境下运行。我们的目标就是在精度下降可控的范围内（比如mAP下降不超过1-2%），尽可能多地“剪掉”冗余，换取速度的成倍提升。

2. 准备工作与环境搭建

在开始动手之前，我们需要把环境和工具准备好。这里我选择使用 Torch Pruning 这个工具库，它基于PyTorch，功能比较全面，而且社区活跃，用起来也相对简单。

首先，确保你已经有了一个训练好的YOLOv12模型权重文件（比如 yolov12s.pt）。我们将在它的基础上进行剪枝。

接下来，安装必要的库。除了PyTorch，我们主要需要 torch-pruning 这个库。

# 安装 torch-pruning
pip install torch-pruning

# 确保你有最新版的 torch 和 torchvision
pip install --upgrade torch torchvision

如果你的YOLOv12代码是基于Ultralytics YOLO的，可能还需要安装 ultralytics 包。这里我们假设你已经有了一个可以正常加载和推理的YOLOv12模型脚本。

为了评估剪枝效果，我们需要一个验证数据集，比如COCO val2017的一部分，用来在剪枝过程中监控模型的精度变化。准备好这些，我们就可以开始了。

3. 理解结构化剪枝：剪什么，怎么剪？

剪枝分为非结构化剪枝和结构化剪枝。非结构化剪枝是细粒度的，它剪掉的是单个权重参数，这会导致模型变得稀疏。虽然压缩率高，但需要特殊的硬件或软件库来加速，通用性不强。

我们这次用的是结构化剪枝，它是粗粒度的，以整个滤波器（Filter）或通道（Channel）为单位进行裁剪。想象一下卷积层，它由很多个滤波器组成，每个滤波器负责提取一种特征。结构化剪枝就是判断哪些滤波器是“不重要”的，然后把它们整个移除。同时，下一层对应的输入通道也要一并移除。

这样做的好处是，剪枝后的模型仍然是一个标准的、密集的神经网络，可以直接被任何支持PyTorch的硬件或框架加载和加速，不需要任何特殊处理，部署起来非常方便。

那么，关键问题来了：怎么判断一个滤波器重不重要？ 常用的准则有：

• L1 Norm（权重绝对值之和）：一个滤波器的所有权重绝对值加起来很小，说明它激活弱，可能不重要。
• BN层缩放因子（Gamma）：如果模型有BatchNorm层，其缩放因子γ的大小可以衡量通道的重要性。γ接近0的通道，可以被剪掉。
• 梯度信息：根据反向传播的梯度大小来判断重要性。

在接下来的实战中，我们会用到基于L1 Norm的准则来剪枝。

4. 实战：一步步剪枝YOLOv12模型

下面我们进入核心环节。我会把代码分成几个关键步骤，并加上详细注释。

4.1 加载预训练模型与数据

首先，加载我们训练好的YOLOv12模型，并准备好验证数据加载器。

import torch
import torch_pruning as tp
from models.yolo import Model  # 假设这是你的YOLOv12模型定义
from utils.dataloaders import create_dataloader  # 假设的数据加载器
from utils.metrics import ap_per_class  # 用于计算mAP

# 1. 加载预训练模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ckpt = torch.load('yolov12s.pt', map_location=device)
model = Model(cfg='models/yolov12s.yaml').to(device)  # 根据你的配置文件
model.load_state_dict(ckpt['model'].float().state_dict(), strict=False)
model.eval()

# 2. 准备验证集（这里以COCO为例，实际请替换为你的数据路径）
val_loader = create_dataloader('path/to/coco/val2017.txt', imgsz=640, batch_size=8, stride=32, pad=0.5, workers=4)[0]

# 3. 定义一个评估函数，用于计算模型在验证集上的mAP
def evaluate_model(model, dataloader):
    # 这里简化了评估流程，实际需要运行完整的检测和mAP计算
    # 假设我们有一个函数 run_val 返回mAP
    # stats = run_val(model, dataloader)
    # return stats[0]  # 返回 mAP@0.5
    print("进行模型评估...")
    # 此处为演示，返回一个模拟的基线mAP
    return 0.45  # 假设基线mAP是0.45

base_map = evaluate_model(model, val_loader)
print(f"基准模型 mAP@0.5: {base_map:.4f}")

4.2 构建依赖图并选择要剪枝的层

剪枝不是随便剪的，因为网络层之间有关联。剪掉某一层的输出通道，下一层对应的输入通道也必须被剪掉。torch-pruning 库会自动帮我们处理这种依赖关系。

我们通常选择卷积层（Conv2d）进行剪枝。对于YOLO这类检测模型，需要谨慎对待靠近输出的层（检测头），因为它们对最终定位和分类影响很大。

# 4. 构建依赖图，分析层间依赖
example_inputs = torch.randn(1, 3, 640, 640).to(device)
DG = tp.DependencyGraph().build_dependency(model, example_inputs=example_inputs)

# 5. 选择要剪枝的层（这里我们选择所有卷积层，除了最后几层）
pruning_layers = []
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        # 可以排除某些关键层，比如检测头的最后一层卷积
        if 'detect' not in name:  # 简单过滤，根据你的模型结构调整
            pruning_layers.append(module)
print(f"总共找到 {len(pruning_layers)} 个可剪枝的卷积层。")

4.3 执行迭代式剪枝

我们不会一次性剪掉太多，而是采用迭代式剪枝：剪一点，评估一下精度，如果精度下降在可接受范围内，就继续剪。

# 6. 定义剪枝函数（基于L1 Norm）
def prune_model(model, pruning_layers, prune_ratio=0.2):
    """
    对指定的层进行剪枝。
    prune_ratio: 每层要剪掉的比例（例如0.2表示剪掉20%的滤波器）
    """
    pruning_plan = []
    for layer in pruning_layers:
        # 获取该层权重的L1 Norm
        weight = layer.weight.data
        l1_norm = weight.abs().sum(dim=(1,2,3)) # 计算每个滤波器的L1 Norm
        num_filters = len(l1_norm)
        num_to_prune = int(num_filters * prune_ratio)
        
        if num_to_prune > 0:
            # 找到L1 Norm最小的那些滤波器的索引
            prune_indices = torch.argsort(l1_norm)[:num_to_prune].tolist()
            # 创建剪枝计划
            pruning_plan.append((layer, tp.prune_conv_out_channel, prune_indices))
    
    # 执行剪枝计划，依赖图会自动处理关联层的剪枝
    for dep in pruning_plan:
        layer, prune_func, indices = dep
        prune_func(layer, indices)
        # 更新依赖图，为下一轮剪枝做准备
        DG.prune(layer, prune_func, idxs=indices)
    
    # 剪枝后，模型结构变了，需要手动“压缩”一下，移除被剪掉的通道
    # torch-pruning 提供了 make_pruned_model 函数，但这里我们演示手动处理依赖后，模型已更新。
    # 实际上，更常用的方式是使用 tp.prune_model 的高级API，下面会展示。
    return model

# 7. 迭代剪枝与评估
target_map_drop = 0.02  # 允许mAP下降的最大值，例如2%
current_map = base_map
prune_ratio_per_iter = 0.05  # 每轮剪枝5%
iteration = 0

while (base_map - current_map) < target_map_drop:
    iteration += 1
    print(f"\n--- 开始第 {iteration} 轮剪枝 ---")
    
    # 使用torch-pruning提供的高级API进行剪枝，更简洁
    # 策略：基于L1 Norm，全局剪枝（考虑所有选中层的整体重要性）
    pruning_strategy = tp.strategy.L1Strategy()
    # 我们依然只剪枝之前选中的那些卷积层
    pruning_idxs = {}
    for i, layer in enumerate(pruning_layers):
        # 计算每层要剪的数量
        num_filters = layer.out_channels
        num_to_prune = int(num_filters * prune_ratio_per_iter)
        if num_to_prune > 0:
            # 使用策略选择要剪的索引
            pruning_idxs[layer] = pruning_strategy(layer.weight, amount=num_to_prune)
    
    # 执行剪枝
    if pruning_idxs:
        # 重要：使用依赖图来安全剪枝
        for layer, idxs in pruning_idxs.items():
            DG.prune(layer, tp.prune_conv_out_channel, idxs=idxs)
        print(f"本轮剪掉了约 {prune_ratio_per_iter*100:.1f}% 的滤波器。")
    else:
        print("本轮没有可剪枝的层，停止迭代。")
        break
    
    # 评估剪枝后的模型
    current_map = evaluate_model(model, val_loader)
    print(f"当前模型 mAP@0.5: {current_map:.4f}, 累计下降: {base_map - current_map:.4f}")
    
    # 如果精度下降接近阈值，减小剪枝比例或停止
    if (base_map - current_map) > target_map_drop * 0.8:
        prune_ratio_per_iter *= 0.5  # 减半剪枝比例
        print(f"精度下降较快，将下轮剪枝比例调整为 {prune_ratio_per_iter:.3f}")

print(f"\n剪枝结束。最终mAP: {current_map:.4f} (基准: {base_map:.4f})")

4.4 微调与保存剪枝后模型

剪枝操作会破坏模型原本学到的特征分布，因此精度通常会有所下降。为了恢复甚至提升精度，微调（Fine-tuning）是必不可少的一步。

# 8. 对剪枝后的模型进行微调
def fine_tune_model(model, train_loader, epochs=10):
    """
    简化的微调训练循环。
    """
    model.train()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    criterion = ...  # 定义你的损失函数，例如YOLO的复合损失
    
    for epoch in range(epochs):
        for batch_i, (imgs, targets, paths, _) in enumerate(train_loader):
            imgs = imgs.to(device)
            targets = targets.to(device)
            
            optimizer.zero_grad()
            outputs = model(imgs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            
        print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}')
        # 可以每隔几个epoch在验证集上评估一次
        if (epoch+1) % 5 == 0:
            model.eval()
            val_map = evaluate_model(model, val_loader)
            print(f'  验证集 mAP@0.5: {val_map:.4f}')
            model.train()
    
    return model

# 加载训练数据（示例）
train_loader = create_dataloader('path/to/coco/train2017.txt', imgsz=640, batch_size=16, stride=32, pad=0.5, workers=8)[0]
print("开始微调剪枝后的模型...")
model = fine_tune_model(model, train_loader, epochs=20)
model.eval()

# 9. 保存剪枝并微调后的模型
final_map = evaluate_model(model, val_loader)
print(f"微调后最终 mAP@0.5: {final_map:.4f}")

pruned_model_path = 'yolov12s_pruned_finetuned.pt'
torch.save({
    'model': model.state_dict(),
    'map': final_map
}, pruned_model_path)
print(f"剪枝后模型已保存至: {pruned_model_path}")

5. 效果对比与结果分析

做完以上所有步骤，我们来对比一下剪枝前后的关键指标。这里我列一个简单的表格，展示理想情况下的效果：

指标	原始YOLOv12模型	剪枝后模型 (目标)	变化幅度
参数量 (Params)	~15M	~9M	减少约40%
计算量 (FLOPs)	~30G	~18G	减少约40%
模型文件大小	~30MB	~18MB	缩小约40%
推理速度 (GPU)	15 ms/img	9 ms/img	提升约40%
推理速度 (CPU)	200 ms/img	120 ms/img	提升约40%
mAP@0.5	0.450	0.445	下降约0.5%

注：以上为示例数据，实际减少比例和精度损失取决于你的模型、数据集和剪枝策略。

从表格可以看出，通过结构化剪枝，我们在精度损失极小（仅0.5%） 的情况下，成功将模型的参数量、计算量和文件大小都减少了约40%。这直接带来了推理速度的显著提升，无论是在GPU还是CPU上，速度提升都接近40%。这意味着在同样的硬件上，你的应用可以处理更多的视频流，或者响应更加实时。

6. 总结与建议

这次带大家走了一遍YOLOv12模型剪枝的完整流程。整个过程就像给一棵树做修剪，目的是让它长得更健康、更高效，而不是破坏它。核心就是利用工具自动分析网络各层的重要性，然后安全、迭代地移除冗余部分。

实际操作下来，有几点感受比较深。首先，依赖图这个工具真的很重要，它能避免我们手动剪枝时出错，确保网络结构在剪枝后依然是正确连通的。其次，迭代剪枝比一次性猛剪要稳妥得多，边剪边评估，能把精度损失控制在预期范围内。最后，微调绝不是可有可无的步骤，剪枝后的模型就像动了一场手术，需要一段时间的“康复训练”才能恢复到最佳状态，甚至有可能通过微调学到更泛化的特征。

如果你也想尝试给自己的模型瘦身，我的建议是：

1. 从预训练模型开始：在一个已经收敛的模型上剪枝，效果比从头训练一个稀疏模型要好。
2. 明确你的目标：是追求极致的速度，还是平衡速度和精度？这决定了你剪枝的激进程度。
3. 准备好评估数据：一个具有代表性的验证集对于监控剪枝过程中的精度变化至关重要。
4. 先主干，后头部：对检测模型，可以更激进地剪枝骨干网络（Backbone），而对检测头（Head）要更加保守。
5. 结合其他优化技术：剪枝可以和量化（Quantization）、知识蒸馏（Knowledge Distillation） 等技术结合使用，获得更大的压缩和加速收益。

模型剪枝是一门实践性很强的技术，不同的模型、不同的数据集，最优的剪枝策略可能都不一样。多动手实验，观察分析剪枝前后模型的变化，你就能越来越得心应手，让手中的AI模型在边缘端“飞”起来。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。