BEV多任务知识蒸馏实战：如何解决自动驾驶模型效率与精度平衡难题

1次阅读

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

自动驾驶中的 BEV（Bird’s Eye View）多任务模型通常需要同时处理目标检测、语义分割和轨迹预测等任务。这种多任务联合训练的方式虽然能共享特征提取器，减少计算冗余，但也带来了显著的显存瓶颈：

显存消耗大：多任务模型的参数量往往是单任务模型的数倍，尤其是 BEV 特征图的高分辨率特性，使得显存占用成倍增加。
任务间干扰：不同任务对特征的需求可能存在冲突，例如目标检测需要明确的物体边界，而语义分割则需要连续的语义信息。
传统压缩方法的局限性：
剪枝可能会破坏多任务共享的特征结构
量化在多任务场景下难以平衡不同任务的精度需求
知识蒸馏成为更灵活的选择

特征蒸馏：直接对齐教师和学生网络的中间层特征
优点：实现简单
缺点：对不同任务的适应性差
关系蒸馏：捕捉任务间的相关性
公式：$L_{rel} = \sum_{i,j}(S_{ij}-T_{ij})^2$
其中 $S_{ij}$ 和 $T_{ij}$ 分别表示学生和教师网络的特征相似度
逻辑蒸馏：在预测层面进行蒸馏
更适合处理多任务输出的不一致性问题

我们设计了一种新型的跨任务注意力蒸馏（CTAD）模块：

class CTAD(nn.Module):
    def __init__(self, num_tasks):
        super().__init__()
        self.task_projections = nn.ModuleList([nn.Conv2d(256, 64, 1) for _ in range(num_tasks)
        ])

    def forward(self, teacher_feats, student_feats):
        # 计算任务间注意力权重
        attn_maps = []
        for proj in self.task_projections:
            attn = torch.sigmoid(proj(teacher_feats))
            attn_maps.append(attn)
        # [...] 完整实现省略

使用 MMDetection3D 框架构建基线模型：

教师网络：基于 Swin-Transformer 的大型 BEV 多任务模型
学生网络：精简版的 CNN+Transformer 混合结构

多尺度特征对齐模块：

def feature_align(teacher_feat, student_feat):
    # 使用可变形卷积进行特征对齐
    align_conv = DeformConv2d(teacher_feat.size(1), student_feat.size(1), 3)
    aligned_feat = align_conv(teacher_feat)
    return F.mse_loss(aligned_feat, student_feat)

任务相关性权重计算：

def compute_task_weights(losses):
    # 基于各任务 loss 的指数移动平均动态调整权重
    weights = [1.0 / (l.mean().detach() + 1e-6) for l in losses]
    return F.softmax(torch.stack(weights), dim=0)

在 nuScenes 数据集上的测试结果：