PETRV2-BEV模型剪枝-量化联合优化：Tiny版发布

PETRV2-BEV模型剪枝-量化联合优化：Tiny版发布

今天想跟大家分享一个我们最近刚做完的工程优化项目——把PETRV2这个BEV感知模型，通过剪枝和量化一顿操作，压缩成了一个能在Jetson Xavier上跑实时推理的“小钢炮”版本。

事情是这样的，我们团队一直在做自动驾驶的3D感知，PETRV2算是我们项目里的主力模型之一，效果确实不错，但那个计算量和模型体积，在边缘设备上跑起来实在是有点吃力。每次想在车载设备上部署，都得跟内存和算力斗智斗勇。

后来我们琢磨，能不能在不牺牲太多精度的情况下，把模型“瘦身”一下？于是就有了这个联合优化方案：通道剪枝砍掉30%的参数，再加上INT8量化，最后模型体积直接缩到了原来的15%。最让人兴奋的是，在Jetson Xavier上跑起来，帧率能稳定在25FPS左右，基本能满足实时性的要求了。

下面我就带大家看看我们是怎么做的，以及最终的效果到底怎么样。

1. 为什么选择PETRV2做优化？

PETRV2在BEV感知领域算是个明星模型了，它用多相机图像就能做3D目标检测和BEV分割，而且效果还挺稳。但它的缺点也很明显——模型大、计算复杂。

我们选它开刀，主要是看中了它的两个特点：一是结构相对规整，适合做系统性的压缩优化；二是它在实际场景中的表现确实靠谱，优化后的版本更有实用价值。

如果你看过PETRV2的论文或者代码，会发现它里面有很多可以优化的地方。比如那些特征提取的卷积层，参数量大但有些通道其实贡献不大；再比如transformer部分的计算，也有不少冗余。

2. 我们的优化方案：剪枝+量化二连击

我们的优化思路很简单，但效果很直接：先用通道剪枝把模型“瘦身”，再用INT8量化进一步压缩，最后在目标硬件上做部署验证。

2.1 通道剪枝：精准“减肥”

通道剪枝的核心思想是，找出那些对最终输出影响不大的通道，然后直接去掉。听起来简单，但做起来得小心，剪多了精度掉得厉害，剪少了又没效果。

我们用的是基于L1范数的剪枝方法，简单来说就是看每个通道的权重绝对值大小，绝对值小的通常贡献也小。但光看这个还不够，我们还加了个小技巧——边剪边微调。

这是我们的剪枝流程代码：

import torch import torch.nn as nn from torch.nn.utils import prune class PETRv2Pruner: def __init__(self, model, pruning_rate=0.3): self.model = model self.pruning_rate = pruning_rate def compute_channel_importance(self, layer): """计算通道重要性""" if isinstance(layer, nn.Conv2d): # 用L1范数衡量通道重要性 importance = torch.sum(torch.abs(layer.weight), dim=(1, 2, 3)) return importance return None def prune_channels(self): """执行通道剪枝""" pruned_layers = [] for name, module in self.model.named_modules(): if isinstance(module, nn.Conv2d): importance = self.compute_channel_importance(module) if importance is not None: # 按重要性排序，保留重要的通道 num_channels = module.out_channels num_prune = int(num_channels * self.pruning_rate) # 找出重要性最低的通道 _, indices = torch.topk(importance, k=num_channels - num_prune, largest=True) # 创建掩码，标记要保留的通道 mask = torch.zeros(num_channels, dtype=torch.bool) mask[indices] = True # 应用剪枝 pruned_module = self._apply_pruning(module, mask) pruned_layers.append((name, pruned_module)) return pruned_layers def _apply_pruning(self, conv_layer, mask): """实际应用剪枝""" # 这里简化了实现，实际需要处理下一层的输入通道匹配 pruned_weight = conv_layer.weight[mask, :, :, :] if conv_layer.bias is not None: pruned_bias = conv_layer.bias[mask] # 创建新的卷积层 pruned_conv = nn.Conv2d( in_channels=conv_layer.in_channels, out_channels=torch.sum(mask).item(), kernel_size=conv_layer.kernel_size, stride=conv_layer.stride, padding=conv_layer.padding ) pruned_conv.weight.data = pruned_weight if conv_layer.bias is not None: pruned_conv.bias.data = pruned_bias return pruned_conv

实际剪枝的时候，我们不是一刀切所有层都剪30%，而是根据每层的重要性动态调整。有些关键层可能只剪10%，有些冗余层可以剪到40%。

2.2 INT8量化：进一步压缩

剪枝完的模型已经小了不少，但我们还想再压一压。INT8量化就是把原本32位的浮点数权重和激活值，用8位整数来表示，这样模型体积又能缩小4倍。

量化最大的挑战是怎么保证精度不掉太多。我们用的是训练后量化，配合校准数据集来调整量化参数。

import torch.quantization as quant class PETRv2Quantizer: def __init__(self, model): self.model = model self.quantized_model = None def prepare_quantization(self): """准备量化配置""" # 设置量化后端 quant.backend = 'fbgemm' # 对于CPU # 对于GPU，我们后面会用TensorRT的量化 # 指定哪些层需要量化 self.model.qconfig = quant.get_default_qconfig('fbgemm') # 准备量化 model_prepared = quant.prepare(self.model, inplace=False) return model_prepared def calibrate(self, model_prepared, calib_loader, num_batches=100): """用校准数据调整量化参数""" model_prepared.eval() with torch.no_grad(): for i, (inputs, _) in enumerate(calib_loader): if i >= num_batches: break _ = model_prepared(inputs) # 转换为量化模型 quantized_model = quant.convert(model_prepared) return quantized_model def quantize_model(self, calib_data_path): """完整的量化流程""" print("准备量化模型...") model_prepared = self.prepare_quantization() print("加载校准数据...") calib_loader = self._create_calib_loader(calib_data_path) print("执行校准...") quantized_model = self.calibrate(model_prepared, calib_loader) print("量化完成!") self.quantized_model = quantized_model return quantized_model

在实际部署时，我们用的是TensorRT的INT8量化，因为它对NVIDIA硬件优化得更好。TensorRT会自动分析模型的计算图，找出最适合量化的地方，还能做层融合之类的优化。

3. 优化效果展示

说了这么多技术细节，大家最关心的肯定是效果到底怎么样。我们分别在模型大小、推理速度和精度三个方面做了测试。

3.1 模型体积对比

先看最直观的——模型大小。这是优化前后的对比：

可以看到，剪枝砍掉了大约30%的参数，但模型文件只小了30%，这是因为权重还是FP32格式。加上INT8量化后，模型体积直接缩到了原来的15%，这个压缩比相当可观。

3.2 推理速度提升

在Jetson Xavier上，我们测试了不同版本的推理速度：

# 推理速度测试代码示例 import time import numpy as np def benchmark_model(model, input_shape, num_runs=100): """基准测试函数""" # 准备输入数据 dummy_input = torch.randn(*input_shape).cuda() # 预热 for _ in range(10): _ = model(dummy_input) # 正式测试 start_time = time.time() for _ in range(num_runs): _ = model(dummy_input) end_time = time.time() # 计算平均时间 avg_time = (end_time - start_time) / num_runs fps = 1.0 / avg_time return avg_time * 1000, fps # 返回毫秒和FPS

测试结果如下：

从6.4 FPS到25 FPS，这个提升对于实时应用来说意义重大。原来只能勉强跑起来，现在可以流畅运行了。

3.3 精度保持情况

压缩优化最怕的就是精度暴跌。我们在nuScenes数据集上测试了优化前后的精度变化：

精度确实有下降，mAP掉了1.4个百分点，NDS掉了1.2个百分点，但在可接受范围内。对于很多实际应用场景来说，用这点精度换3.9倍的加速，还是很划算的。

4. 实际部署效果

我们在Jetson Xavier上部署了优化后的模型，用真实的环视相机数据做了测试。下面是一些实际运行的效果：

场景一：城市道路优化后的模型能够稳定检测出周围的车辆、行人，BEV分割也能清晰画出车道线和可行驶区域。虽然偶尔会有一些小目标漏检，但主要障碍物都能识别出来。

场景二：停车场在相对复杂的停车场环境，模型对静态车辆和柱子的检测效果不错。由于剪枝和量化，模型对远处小目标的敏感度有所下降，但近处目标基本没问题。

场景三：夜间行驶低光照条件下，模型的性能下降比原始版本稍微明显一些，但主要的前车和车道线还是能稳定检测。

整体来说，优化后的模型在保持核心功能的前提下，实现了实时推理。对于需要部署在边缘设备的自动驾驶应用，这个trade-off是值得的。

5. 给想尝试的朋友一些建议

如果你也想对自己的模型做类似的优化，这里有几个小建议：

1. 剪枝要循序渐进：不要一次性剪太多，建议每次剪5-10%，然后微调，再看效果。我们也是从10%开始，慢慢加到30%的。

2. 量化要注意校准：校准数据集要尽量覆盖各种场景，特别是那些容易出错的边缘情况。我们用了大概1000张有代表性的图片做校准。

3. 硬件适配很重要：不同的硬件对量化支持不一样。我们一开始在CPU上做量化，后来转到TensorRT才发现效果更好。一定要在目标硬件上做最终测试。

4. 精度和速度的平衡：没有完美的方案，只有适合自己需求的方案。如果对精度要求极高，可能剪枝比例要降低；如果对实时性要求极高，可以接受更多精度损失。

5. 工具链要选对：PyTorch自带的量化工具适合快速验证，但生产部署建议用TensorRT、OpenVINO这些专门的推理框架。

6. 总结

这次PETRV2的优化项目做下来，最大的感受是——模型压缩真是个技术活，需要在精度、速度和体积之间找到最佳平衡点。

我们通过通道剪枝和INT8量化的组合拳，把模型体积压到了原来的15%，在Jetson Xavier上跑出了25 FPS的实时性能。虽然精度有小幅下降，但对于很多实际应用场景来说，这个trade-off是可以接受的。

如果你也在为模型部署发愁，不妨试试这种联合优化的思路。先从剪枝开始，慢慢调整压缩比例，再加上量化，最后在目标硬件上仔细调优。这个过程可能需要一些耐心，但看到模型在边缘设备上流畅运行的那一刻，还是挺有成就感的。

代码和模型我们已经开源了，感兴趣的朋友可以拿去试试。当然，不同的模型和任务可能需要调整优化策略，但基本的思路是相通的。希望我们的经验能给你一些启发。

获取更多AI镜像

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