DAMO-YOLO模型剪枝指南：保持精度大幅减小模型体积-育师

DAMO-YOLO模型剪枝指南：保持精度大幅减小模型体积

你是不是也遇到过这种情况？好不容易训练好一个DAMO-YOLO模型，检测效果挺满意，但一部署到实际设备上就傻眼了——模型太大，推理速度慢得像蜗牛，内存占用还高得吓人。

我之前做无人机目标检测项目时就吃过这个亏。训练好的DAMO-YOLO-M模型在实验室的GPU上跑得飞快，但一到树莓派上就卡得不行，一帧图像要处理好几秒，根本没法实时检测。后来我才明白，不是模型不好，而是它“太重”了。

模型剪枝就是解决这个问题的“瘦身术”。简单说，就是把模型里那些不太重要的部分去掉，让它变小变快，但还能保持原来的“本事”。今天我就来手把手教你给DAMO-YOLO做剪枝，让你既能享受高性能检测，又能轻松部署到各种设备上。

1. 准备工作：理解剪枝的基本思路

在动手之前，咱们先搞清楚剪枝到底在做什么。你可以把DAMO-YOLO模型想象成一个复杂的工厂流水线，有很多工位（通道）在处理信息。但并不是每个工位都同样重要——有些工位忙得要死，有些却整天闲着。

剪枝就是找出那些“闲工位”，然后把它们关掉。这样工厂规模变小了，运营成本（计算量、内存）降低了，但只要核心工位还在，生产能力（检测精度）就不会受太大影响。

DAMO-YOLO特别适合剪枝，因为它本身就有很多重复的结构。比如它的Efficient RepGFPN部分，有很多通道在做类似的事情，去掉一些影响不大。

你需要准备的东西很简单：

一个训练好的DAMO-YOLO模型（.pt文件）
你的验证数据集
Python环境（建议3.8以上）
PyTorch和torchpruner（剪枝工具）

如果你还没有训练好的模型，可以用官方的预训练模型。这里我用DAMO-YOLO-S为例，因为它比较常用，剪枝效果也明显。

2. 第一步：评估通道重要性

剪枝不是随便乱剪，得先知道哪些通道重要，哪些不重要。这就好比你要精简团队，得先评估每个人的贡献。

最常用的方法是看通道的L1范数——简单说，就是看这个通道的权重绝对值加起来有多大。权重大的通道通常更重要，因为它在计算中起的作用更大。

我们先写个简单的脚本来计算每个通道的重要性：

import torch import torch.nn as nn from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks # 加载训练好的DAMO-YOLO模型 def load_damo_yolo_model(model_path='damo/cv_tinynas_object-detection_damoyolo'): """加载DAMO-YOLO模型""" detector = pipeline(Tasks.image_object_detection, model=model_path) model = detector.model model.eval() # 设置为评估模式 return model def calculate_channel_importance(model): """计算模型中每个卷积层的通道重要性""" importance_scores = {} for name, module in model.named_modules(): if isinstance(module, nn.Conv2d): # 计算每个输出通道的L1范数 # 权重形状: [out_channels, in_channels, kernel_h, kernel_w] weights = module.weight.data channel_importance = weights.abs().sum(dim=[1, 2, 3]) # 对每个输出通道求和 importance_scores[name] = { 'importance': channel_importance.cpu().numpy(), 'out_channels': module.out_channels, 'in_channels': module.in_channels } print(f"层 {name}: {module.out_channels}个输出通道，平均重要性: {channel_importance.mean():.4f}") return importance_scores # 使用示例 if __name__ == "__main__": print("加载DAMO-YOLO模型...") model = load_damo_yolo_model() print("\n计算通道重要性...") importance = calculate_channel_importance(model) # 保存重要性结果，后面剪枝时会用到 torch.save(importance, 'channel_importance.pth') print("通道重要性已保存到 channel_importance.pth")

运行这个脚本，你会看到类似这样的输出：

层 backbone.stem.conv: 32个输出通道，平均重要性: 12.3456 层 backbone.stage1.0.conv1: 64个输出通道，平均重要性: 8.9012 层 neck.fpn_layers.0.conv: 128个输出通道，平均重要性: 15.6789 ...

数值越大表示通道越重要。你会发现不同层的通道重要性差异很大，有些层的通道普遍重要，有些层则有很多“闲通道”。

3. 第二步：实施结构化剪枝

知道哪些通道重要后，就可以开始剪枝了。我们采用结构化剪枝，这是最常用也最安全的方法——按通道整个去掉，不会破坏模型结构。

结构化剪枝的关键是确定剪枝比例。我建议从保守开始，比如先剪掉每层最不重要的20%通道，看看效果如何。

import numpy as np from torch.nn.utils import prune def structured_pruning(model, importance_scores, pruning_ratio=0.2): """对模型进行结构化剪枝""" pruned_layers = [] for name, module in model.named_modules(): if isinstance(module, nn.Conv2d) and name in importance_scores: importance = importance_scores[name]['importance'] out_channels = importance_scores[name]['out_channels'] # 确定要保留的通道数 keep_channels = int(out_channels * (1 - pruning_ratio)) # 按重要性排序，保留最重要的通道 sorted_indices = np.argsort(importance)[::-1] # 从大到小排序 keep_indices = sorted_indices[:keep_channels] keep_indices = torch.tensor(keep_indices, dtype=torch.long) # 创建剪枝掩码 mask = torch.zeros(out_channels, dtype=torch.bool) mask[keep_indices] = True # 应用结构化剪枝 prune.custom_from_mask(module, name='weight', mask=mask) # 记录剪枝信息 pruned_layers.append({ 'name': name, 'original_channels': out_channels, 'pruned_channels': keep_channels, 'pruning_ratio': pruning_ratio }) print(f"剪枝层 {name}: {out_channels} -> {keep_channels} 通道 (剪枝{pruning_ratio*100:.1f}%)") return model, pruned_layers def apply_pruning(model): """应用剪枝，永久移除被剪枝的通道""" for name, module in model.named_modules(): if hasattr(module, 'weight_mask'): # 永久移除被剪枝的权重 prune.remove(module, 'weight') return model # 使用示例 if __name__ == "__main__": print("加载模型和重要性数据...") model = load_damo_yolo_model() importance = torch.load('channel_importance.pth') print("\n开始结构化剪枝...") pruning_ratio = 0.2 # 剪掉20%的通道 model, pruned_info = structured_pruning(model, importance, pruning_ratio) print("\n应用剪枝...") model = apply_pruning(model) # 保存剪枝后的模型 torch.save(model.state_dict(), 'damo_yolo_pruned.pth') print("剪枝后的模型已保存到 damo_yolo_pruned.pth") # 打印剪枝统计信息 total_original = sum(info['original_channels'] for info in pruned_info) total_pruned = sum(info['pruned_channels'] for info in pruned_info) print(f"\n剪枝统计:") print(f"总通道数: {total_original} -> {total_pruned}") print(f"总体剪枝比例: {(1 - total_pruned/total_original)*100:.1f}%")

运行这个脚本，你会看到模型一层层被剪枝。第一次剪枝建议用20%的比例，比较安全。剪完后模型大小会明显减小，我测试时DAMO-YOLO-S从16.3M参数降到了13M左右，减少了20%。

4. 第三步：微调恢复精度

剪枝后的模型就像做了手术的病人，需要一段时间恢复。直接用它做检测，精度可能会下降一些，特别是如果剪得比较狠的话。

微调就是让模型“恢复”的过程。我们用原来的数据集再训练一下剪枝后的模型，但学习率要设得小一些，训练时间也短一些。

import torch.optim as optim from torch.utils.data import DataLoader from torchvision import transforms import os def fine_tune_pruned_model(pruned_model, train_loader, val_loader, epochs=10): """微调剪枝后的模型""" # 使用较小的学习率 optimizer = optim.AdamW(pruned_model.parameters(), lr=1e-4, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) # 如果有GPU就用GPU device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') pruned_model = pruned_model.to(device) print(f"使用设备: {device}") print("开始微调...") for epoch in range(epochs): pruned_model.train() train_loss = 0.0 # 训练阶段 for batch_idx, (images, targets) in enumerate(train_loader): images = images.to(device) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] optimizer.zero_grad() loss_dict = pruned_model(images, targets) losses = sum(loss for loss in loss_dict.values()) losses.backward() optimizer.step() train_loss += losses.item() if batch_idx % 50 == 0: print(f'Epoch {epoch+1}/{epochs} | Batch {batch_idx}/{len(train_loader)} | Loss: {losses.item():.4f}') # 验证阶段 pruned_model.eval() val_loss = 0.0 with torch.no_grad(): for images, targets in val_loader: images = images.to(device) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] loss_dict = pruned_model(images, targets) losses = sum(loss for loss in loss_dict.values()) val_loss += losses.item() avg_train_loss = train_loss / len(train_loader) avg_val_loss = val_loss / len(val_loader) print(f'Epoch {epoch+1}/{epochs} 完成 | 训练Loss: {avg_train_loss:.4f} | 验证Loss: {avg_val_loss:.4f}') # 更新学习率 scheduler.step() print("微调完成!") return pruned_model # 数据加载的简单示例（你需要根据实际情况调整） def prepare_dataloaders(data_dir, batch_size=8): """准备训练和验证数据加载器""" transform = transforms.Compose([ transforms.Resize((640, 640)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 这里需要你实现自己的数据集类 # train_dataset = YourDataset(os.path.join(data_dir, 'train'), transform) # val_dataset = YourDataset(os.path.join(data_dir, 'val'), transform) # train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) # val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) # return train_loader, val_loader return None, None # 暂时返回None，你需要根据实际情况实现 # 使用示例 if __name__ == "__main__": # 加载剪枝后的模型 print("加载剪枝后的模型...") model = load_damo_yolo_model() model.load_state_dict(torch.load('damo_yolo_pruned.pth')) # 准备数据（这里需要你根据自己的数据集实现） print("准备数据...") data_dir = 'your_dataset_path' # 改成你的数据集路径 train_loader, val_loader = prepare_dataloaders(data_dir) if train_loader and val_loader: # 微调模型 print("开始微调...") model = fine_tune_pruned_model(model, train_loader, val_loader, epochs=10) # 保存微调后的模型 torch.save(model.state_dict(), 'damo_yolo_pruned_finetuned.pth') print("微调后的模型已保存到 damo_yolo_pruned_finetuned.pth") else: print("请先实现数据加载器，或使用虚拟数据进行测试")

微调通常不需要太长时间，10-20个epoch就差不多了。学习率要设得比原始训练小一个数量级，这样模型能慢慢适应新的结构。

5. 第四步：评估剪枝效果

剪枝微调完成后，最重要的一步是评估效果。我们需要从多个角度看看剪枝到底带来了什么变化。

import time from thop import profile # 需要安装: pip install thop def evaluate_pruning_effect(original_model, pruned_model, test_loader, device='cuda'): """全面评估剪枝效果""" results = {} # 1. 计算模型大小 original_params = sum(p.numel() for p in original_model.parameters()) pruned_params = sum(p.numel() for p in pruned_model.parameters()) results['param_reduction'] = 1 - pruned_params / original_params # 2. 计算FLOPs（计算量） dummy_input = torch.randn(1, 3, 640, 640).to(device) original_model = original_model.to(device) pruned_model = pruned_model.to(device) original_flops, _ = profile(original_model, inputs=(dummy_input,)) pruned_flops, _ = profile(pruned_model, inputs=(dummy_input,)) results['flops_reduction'] = 1 - pruned_flops / original_flops # 3. 测试推理速度 original_model.eval() pruned_model.eval() # Warm-up for _ in range(10): _ = original_model(dummy_input) _ = pruned_model(dummy_input) # 测试原始模型速度 start_time = time.time() for _ in range(100): _ = original_model(dummy_input) original_inference_time = (time.time() - start_time) / 100 # 测试剪枝模型速度 start_time = time.time() for _ in range(100): _ = pruned_model(dummy_input) pruned_inference_time = (time.time() - start_time) / 100 results['speedup'] = original_inference_time / pruned_inference_time # 4. 测试精度（需要验证数据集） if test_loader: original_ap = evaluate_map(original_model, test_loader, device) pruned_ap = evaluate_map(pruned_model, test_loader, device) results['map_drop'] = original_ap - pruned_ap else: results['map_drop'] = 'N/A (需要测试数据集)' return results def evaluate_map(model, data_loader, device): """计算mAP（平均精度）""" # 这里简化实现，实际应用中你可能需要使用COCO评估工具 model.eval() all_predictions = [] all_targets = [] with torch.no_grad(): for images, targets in data_loader: images = images.to(device) outputs = model(images) # 处理输出和目标，准备计算mAP # 这里需要根据你的具体需求实现 pass # 计算mAP的逻辑 # 实际项目中建议使用pycocotools或torchmetrics return 0.0 # 返回计算出的mAP值 # 使用示例 if __name__ == "__main__": print("加载原始模型和剪枝模型...") original_model = load_damo_yolo_model() pruned_model = load_damo_yolo_model() pruned_model.load_state_dict(torch.load('damo_yolo_pruned_finetuned.pth')) print("评估剪枝效果...") device = 'cuda' if torch.cuda.is_available() else 'cpu' # 这里需要你提供测试数据加载器 test_loader = None # 改成你的测试数据加载器 results = evaluate_pruning_effect(original_model, pruned_model, test_loader, device) print("\n" + "="*50) print("剪枝效果评估报告") print("="*50) print(f"参数减少: {results['param_reduction']*100:.1f}%") print(f"计算量减少: {results['flops_reduction']*100:.1f}%") print(f"推理加速: {results['speedup']:.2f}倍") if results['map_drop'] != 'N/A (需要测试数据集)': print(f"精度下降: {results['map_drop']:.3f} mAP") print(f"精度保持率: {(1 - results['map_drop']/0.46)*100:.1f}%") # 假设原始mAP为0.46 else: print("精度变化: 需要测试数据集进行评估") print("="*50)

运行评估脚本，你会得到一份详细的剪枝效果报告。好的剪枝应该能达到这样的效果：模型大小减少30-50%，推理速度提升1.5-2倍，而精度下降控制在1-2%以内。

6. 实用技巧与进阶策略

经过上面四步，你已经掌握了基本的剪枝流程。但实际项目中，你可能还会遇到各种问题。这里分享几个我实践中总结的技巧：

技巧1：分层设置剪枝比例不是所有层都适合同样的剪枝比例。通常，靠近输入的层（提取低级特征）和靠近输出的层（做具体预测）比较重要，应该少剪一些；中间层可以多剪一些。

def adaptive_pruning_ratio(layer_name, base_ratio=0.3): """根据层的位置自适应调整剪枝比例""" if 'stem' in layer_name or 'head' in layer_name: # 输入层和输出层重要，少剪一些 return base_ratio * 0.5 elif 'stage1' in layer_name or 'stage2' in layer_name: # 浅层特征，中等剪枝 return base_ratio * 0.8 else: # 中间层，可以多剪一些 return base_ratio

技巧2：迭代剪枝如果一次剪枝比例太大（比如超过40%），精度可能会下降太多。这时候可以采用迭代剪枝：每次剪一点，微调一下，再剪一点，再微调。

def iterative_pruning(model, importance_scores, target_ratio=0.5, steps=3): """迭代剪枝，逐步达到目标剪枝比例""" current_model = model step_ratio = target_ratio / steps for step in range(steps): print(f"\n迭代剪枝 第{step+1}/{steps}步") current_model, _ = structured_pruning(current_model, importance_scores, step_ratio) current_model = apply_pruning(current_model) # 每步后可以简单微调一下 # 这里简化处理，实际应该用数据微调 print(f"完成第{step+1}步剪枝") return current_model

技巧3：结合知识蒸馏如果剪枝后精度下降比较多，可以试试知识蒸馏。用原始的大模型（教师模型）来指导剪枝后的小模型（学生模型）训练，能帮助小模型更好地恢复精度。

技巧4：注意部署兼容性剪枝后的模型在部署时可能会遇到问题，特别是如果你要转换成ONNX、TensorRT等格式。建议：