CenterNet模型剪枝与微调实战：从理论到轻量化部署

1次阅读

共计 1918 个字符，预计需要花费 5 分钟才能阅读完成。

CenterNet 作为单阶段目标检测的经典模型，其简洁的 anchor-free 设计和高效的性能使其在工业界广泛应用。但在嵌入式设备部署时，模型的计算力和内存占用成为主要瓶颈。传统量化方案虽然能减少模型大小，但对计算量的优化有限；而剪枝方案能从结构上精简模型，更适用于计算资源受限的场景。

在剪枝过程中，评估通道的重要性是关键步骤。我们采用 L1-norm 作为评估指标，计算每个通道权重的绝对值之和，数值越大的通道通常包含更多重要信息。

BN 层的 γ 系数反映了通道的缩放程度，可以直观地衡量通道的重要性。我们通过设定一个阈值，将 γ 系数低于该阈值的通道剪枝掉。这种结构化剪枝方式保持了模型的整体架构，便于后续的微调和部署。

为了恢复剪枝后的模型性能，我们采用知识蒸馏的方法进行微调。选择性能更强的 CenterNet 作为 teacher 模型，通过 KL 散度损失将 teacher 模型的输出分布传递给学生模型。蒸馏温度一般设置在 1 - 5 之间，具体值需要根据任务调整。

import torch
import torch.nn as nn

def prune_channels(model, threshold=0.01):
    for name, module in model.named_modules():
        if isinstance(module, nn.BatchNorm2d):
            # 获取 γ 系数
            gamma = module.weight.data.abs()
            # 计算掩码
            mask = gamma.gt(threshold).float()
            # 冻结梯度
            module.weight.requires_grad = False
            module.bias.requires_grad = False
            # 应用剪枝
            module.weight.data.mul_(mask)
            module.bias.data.mul_(mask)

class DistillLoss(nn.Module):
    def __init__(self, alpha=0.5, temperature=2):
        super().__init__()
        self.alpha = alpha
        self.temperature = temperature
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')
        self.reg_loss = nn.SmoothL1Loss()
        self.cls_loss = nn.CrossEntropyLoss()

    def forward(self, student_out, teacher_out, targets):
        # 分类损失
        cls_loss = self.cls_loss(student_out['cls'], targets['cls'])
        # 回归损失
        reg_loss = self.reg_loss(student_out['reg'], targets['reg'])
        # 蒸馏损失
        distill_loss = self.kl_loss(F.log_softmax(student_out['cls']/self.temperature, dim=1),
            F.softmax(teacher_out['cls']/self.temperature, dim=1)
        ) * (self.temperature**2)
        # 总损失
        total_loss = (1-self.alpha) * (cls_loss + reg_loss) + self.alpha * distill_loss
        return total_loss

我们在 COCO-val 数据集上进行了实验，结果如下：