从 baseline 到 75.3 mIoU:structure_knowledge_distillation 训练脚本全解析
从 baseline 到 75.3 mIoUstructure_knowledge_distillation 训练脚本全解析【免费下载链接】structure_knowledge_distillationThe official code for the paper Structured Knowledge Distillation for Semantic Segmentation. (CVPR 2019 ORAL) and extension to other tasks.项目地址: https://gitcode.com/gh_mirrors/st/structure_knowledge_distillation语义分割是计算机视觉中的重要任务而知识蒸馏技术能让轻量级学生模型从复杂教师模型中学习到丰富的结构化知识。本文将详细解析structure_knowledge_distillation项目如何通过三种蒸馏策略将 Cityscapes 数据集上的 mIoU 从 69.10 提升到 75.3 的完整训练流程。 项目概览与性能提升structure_knowledge_distillation是一个基于 CVPR 2019 口头报告论文的开源实现专注于密集预测任务的结构化知识蒸馏。该项目通过创新的三层蒸馏架构在语义分割、目标检测和深度估计等多个任务上都取得了显著性能提升。在 Cityscapes 数据集上项目展示了令人印象深刻的性能提升轨迹模型配置平均 mIoU提升幅度Baseline69.10- 像素级蒸馏70.511.41 像素级成对蒸馏71.782.68 完整三层蒸馏75.36.20️ 三层蒸馏架构解析1. 像素级蒸馏Pixel-wise Distillation像素级蒸馏是最基础的蒸馏形式让学生模型学习教师模型在像素级别的输出分布。在utils/criterion.py中CriterionPixelWise类实现了这一功能# 核心实现像素级蒸馏损失计算 softmax_pred_T F.softmax(preds_T[0].permute(0,2,3,1).contiguous().view(-1,C), dim1) loss (torch.sum(-softmax_pred_T * logsoftmax(preds_S[0]))) / W / H通过计算教师和学生输出之间的 KL 散度学生模型能够学习到教师模型的软标签信息而不仅仅是硬标签。2. 成对蒸馏Pair-wise Distillation成对蒸馏是项目的创新点之一它关注特征空间中的结构化关系。在CriterionPairWiseforWholeFeatAfterPool类中项目实现了基于池化后特征的相似度计算# 特征相似度计算 total_w, total_h feat_T.shape[2], feat_T.shape[3] patch_w, patch_h int(total_w*self.scale), int(total_h*self.scale) maxpool nn.MaxPool2d(kernel_size(patch_w, patch_h), stride(patch_w, patch_h)) loss self.criterion(maxpool(feat_S), maxpool(feat_T))这种方法让学生模型学习到教师模型特征图中的空间关系模式而不仅仅是单个像素的预测。3. 整体蒸馏Holistic Distillation整体蒸馏采用对抗学习的方式通过判别器让学生模型生成与教师模型相似的整体特征分布。在networks/kd_model.py中对抗训练的核心逻辑如下# 对抗训练实现 if self.args.ho True: d_out_S self.parallel_D(self.preds_S, parallelargs.parallel) G_loss G_loss args.lambda_d * self.criterion_adv_for_G(d_out_S, d_out_S)Cityscapes 数据集上的语义分割结果对比 训练脚本深度解析核心训练流程项目的训练主脚本train_and_eval.py实现了完整的训练循环数据加载使用 Cityscapes 数据集支持随机镜像和缩放增强模型初始化同时加载教师模型ResNet-101和学生模型ResNet-18多阶段训练逐步引入三种蒸馏损失周期性评估每 2975 个批次进行一次验证集评估关键训练参数配置在run_train_val.sh中项目提供了完整的训练配置# 三层蒸馏全部启用 is_pi_useTrue # 像素级蒸馏 is_pa_useTrue # 成对蒸馏 is_ho_useTrue # 整体蒸馏 # 损失权重配置 lambda_pi10.0 # 像素级损失权重 lambda_pa0.5 # 成对损失权重 lambda_d0.1 # 对抗损失权重学习率调度策略项目采用多项式衰减的学习率调度策略在networks/kd_model.py中实现def lr_poly(self, base_lr, iter, max_iter, power): return base_lr * ((1 - float(iter) / max_iter) ** power)这种调度方式在训练初期提供较高的学习率随着训练进行逐渐衰减有助于模型收敛到更好的局部最优解。 环境配置与快速开始系统要求Python 3.5PyTorch 0.4.1CUDA 支持推荐使用 Anaconda 环境编译步骤项目使用 InPlace-ABN 优化内存使用需要编译 CUDA 扩展cd libs sh build.sh python build.py数据集准备下载 Cityscapes 数据集按照项目结构组织数据目录使用预定义的训练/验证列表文件 性能优化技巧1. 批次大小调整默认批次大小为 8根据 GPU 显存可以适当调整。更大的批次大小通常能带来更稳定的梯度估计。2. 损失权重调优三个蒸馏损失的权重需要根据具体任务调整lambda_pi控制像素级蒸馏强度lambda_pa控制成对蒸馏强度lambda_d控制对抗蒸馏强度3. 学习率策略项目默认使用 0.01 的初始学习率可以根据训练曲线调整。如果训练早期损失震荡可以适当降低学习率。 实际应用建议迁移到其他任务项目的三层蒸馏架构可以轻松迁移到其他密集预测任务目标检测在 FCOS 检测器上mAP 从 30.9 提升到 34.0深度估计在 VNL 深度估计模型上相对误差从 13.5% 降低到 13.0%自定义数据集适配要适配自己的数据集需要修改dataset/datasets.py中的数据加载器dataset/list/目录下的文件列表类别数量和对应的配置文件 调试与问题排查常见问题内存不足减小批次大小或输入图像尺寸训练不稳定调整损失权重或降低学习率性能提升不明显检查教师模型质量或数据预处理监控训练过程项目集成了 TensorBoard 日志可以通过以下方式监控训练损失曲线变化验证集 mIoU 趋势各蒸馏损失分量贡献 实验结果分析消融实验项目的消融实验清晰地展示了每层蒸馏的贡献蒸馏组合mIoU相对提升无蒸馏69.10基准仅像素级70.512.0%像素级成对71.783.9%完整三层75.39.0%类别级性能提升在 Cityscapes 的 19 个类别中不同类别的提升幅度不同大尺度物体如道路、建筑提升显著小尺度物体如交通标志、杆子提升相对较小复杂场景如人群、车辆受益于结构化信息 未来扩展方向1. 多任务蒸馏可以扩展为同时进行分割、检测和深度估计的多任务蒸馏框架。2. 在线蒸馏从静态教师模型扩展到在线教师模型实现师生协同进化。3. 自适应权重根据训练进度动态调整各蒸馏损失的权重。 总结与建议structure_knowledge_distillation项目通过创新的三层蒸馏架构为语义分割任务提供了强大的知识迁移方案。对于实际应用建议从简单开始先启用像素级蒸馏逐步加入成对和整体蒸馏耐心调参损失权重需要根据具体任务仔细调整充分利用硬件合理设置批次大小充分利用 GPU 并行能力监控验证集定期在验证集上评估防止过拟合通过本项目的训练脚本和架构设计你可以在保持模型轻量化的同时显著提升语义分割性能实现从 baseline 到 75.3 mIoU 的跨越式提升【免费下载链接】structure_knowledge_distillationThe official code for the paper Structured Knowledge Distillation for Semantic Segmentation. (CVPR 2019 ORAL) and extension to other tasks.项目地址: https://gitcode.com/gh_mirrors/st/structure_knowledge_distillation创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考