ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

MMSegmentation 中的 K-Net 统一图像分割:原理、源码实现与 ADE20K 配置详解

MMSegmentation 中的 K-Net 统一图像分割:原理、源码实现与 ADE20K 配置详解 MMSegmentation 中的 K-Net 统一图像分割原理、源码实现与 ADE20K 配置详解【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation导读本文围绕 configs/knet/README.md 展开系统讲解 K-NetKernel-based Network这一统一图像分割框架在 OpenMMLab MMSegmentation 中的落地方式先梳理其“一组可学习内核统一语义/实例/全景分割”的核心思想与论文要点再结合 knet_head.py 的源码逐层拆解IterativeDecodeHead、KernelUpdateHead、KernelUpdator的实现细节随后逐字段解析 configs/knet 目录下的 7 份 ADE20K 训练配置最后给出实验结果表格与复现命令。读完本文你将理解 K-Net 的动态内核更新机制为何能替代传统框与 NMS并掌握在 MMSegmentation 中从零训练、评测 K-Net 的完整实操路径。K-Net 是什么用一组可学习内核统一三种分割任务K-Net 全称K-Net: Towards Unified Image Segmentation论文发表于 NeurIPS 2021作者为 Wenwei Zhang、Jiangmiao Pang、Kai Chen、Chen Change Loy。其核心观察是语义分割、实例分割与全景分割虽然在传统上分别采用不同的专用框架但它们底层共享“为每个对象/类别生成掩码”这一任务本质。K-Net 用一个统一、简单且有效的框架同时处理这些本质上相似的任务。框架的核心是一组可学习内核learnable kernels每个内核负责为“一个潜在实例”或“一个 stuff 类别”生成一张掩码从而让实例与语义类别被一致地分割出来。为了克服“区分不同实例”的难点论文提出内核更新策略kernel update strategy使每个内核能够动态地、条件地关注输入图像中与其对应的有意义的像素组group实现“内核随图像内容而演化”。由此K-Net 支持端到端训练配合二部匹配 bipartite matching训练与推理天然无需 NMS、无需边界框NMS-free and box-free。论文在原文摘要中报告K-Net 在 MS COCO test-dev 的全景分割上达到 55.2% PQ在 ADE20K val 的语义分割上达到 54.3% mIoU超越了此前所有已发表的最佳单模型结果其实例分割性能与 Cascade Mask R-CNN 相当而推理速度可快 60%–90%。需要说明的是这些数字是论文原文的声明本仓库中的 ADE20K 复现结果以 configs/knet/README.md 与 configs/knet/metafile.yaml 记录的实际评测为准详见下文“ADE20K 实验结果”一节。MMSegmentation 中的实现三个核心模块K-Net 在 MMSegmentation 中完整实现在 mmseg/models/decode_heads/knet_head.py并由 mmseg/models/decode_heads/init.py 统一导出IterativeDecodeHead、KernelUpdateHead、KernelUpdator三个组件注册到模型注册表中供配置直接引用。1. IterativeDecodeHead迭代解码的顶层编排者IterativeDecodeHead源码 knet_head.py#L396-L460是 K-Net 在语义分割中的总入口继承自BaseDecodeHead。其构造函数接收三组关键配置num_stages迭代阶段数即串联的 Kernel Update Head 个数配置中固定为 3kernel_generate_head内核生成头负责产出初始掩码预测、动态内核与类别预测供后续各阶段使用kernel_update_head内核更新头列表逐阶段迭代地精化动态内核与掩码预测。从源码看IterativeDecodeHead.forward()knet_head.py#L433-L450的调用链非常清晰先用kernel_generate_head._forward_feature(inputs)提取融合特征feats再调用其cls_seg得到初始语义分割sem_seg把内核生成头的卷积权重conv_seg.weight克隆并扩展为批维度作为初始seg_kernels依次把feats、seg_kernels、sem_seg送入每一个KernelUpdateHead逐步得到更精细的sem_seg与seg_kernels训练时返回所有阶段的预测stage_segs用于逐阶段监督测试时仅返回最后阶段的预测stage_segs[-1]。对应的loss_by_featknet_head.py#L452-L460会对每个阶段的输出分别计算损失并以loss.s0 / loss.s1 / loss.s2这样的键名汇总实现对“多阶段逐步精化”的显式监督。2. KernelUpdateHead动态内核的迭代精化单元KernelUpdateHead源码 knet_head.py#L145-L393是逐阶段执行内核更新的核心单元。其内部结构包括多头注意力MultiheadAttention作用于展平后的内核向量及其后的 LayerNorm由kernel_updator_cfg指定的内核更新器默认KernelUpdator可选的特征变换卷积feat_transformConvModule与 FFNfeedforward_channels、num_ffn_fcs控制其宽度与层数一组掩码全连接层mask_fcs数量由num_mask_fcs控制与最终输出层fc_mask。其forwardknet_head.py#L282-L393按论文公式推进Group Feature Assembling对应论文 Eq.(3)把上一阶段的掩码做 softmax 后作为权重通过torch.einsum(bnhw,bchw-bnc, sigmoid_masks, x)把特征图聚合成每个组/内核的特征向量源码注释说明 einsum 比 bmm 快约 30%内核更新由KernelUpdator结合聚合特征与当前内核特征产出更新后的内核注意力 FFN对更新后的内核做自注意力与 FFN 变换掩码生成把精化后的内核重塑为卷积核与特征图做F.conv2d分组卷积得到新掩码。源码注释对比了unfold与分组卷积两种实现——分组卷积比 unfold 快约 5 倍且显存约为其 1/52.9ms vs 13.5ms278MB vs 1420MB但实际训练中concat batch方式更快故最终保留了后者。3. KernelUpdator动态卷积式内核更新器KernelUpdator源码 knet_head.py#L18-L142实现论文 Eq.(4)、Eq.(5) 中的动态内核更新机制。它通过两组线性层dynamic_layer对应式中的 φ1/ψ3input_layer对应 φ2/ψ4分别从“更新特征”和“输入特征”生成参数再经门控机制融合# 门控融合对应论文 Eq.(5) features update_gate * param_out.unsqueeze(-2) input_gate * input_out其中gate_sigmoidTrue时对两个门控做 sigmoid之后经过fc_layer、LayerNorm 与激活函数输出新内核。这一“门控 归一化 全连接”的设计让每个内核的更新量取决于其对应像素组的内容正是 K-Net 能够区分不同实例的关键所在。配置文件逐段解析以 ADE20K 为例configs/knet 目录共包含 7 份可训练配置与 metafile.yaml 元数据文件。它们共享同一套骨架我们以 knet-s3_r50-d8_deeplabv3_8xb2-adamw-80k_ade20k-512x512.py 为例完整解读。基础继承与数据预处理_base_ [ ../_base_/datasets/ade20k.py, ../_base_/default_runtime.py, ../_base_/schedules/schedule_80k.py ] crop_size (512, 512) data_preprocessor dict( typeSegDataPreProcessor, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375], bgr_to_rgbTrue, pad_val0, sizecrop_size, seg_pad_val255)数据集基座是 configs/base/datasets/ade20k.py使用ADE20KDataset训练/验证前缀分别为images/training、images/validation训练管线包含RandomResizescale 2048×512ratio_range 0.5–2.0、RandomCropcat_max_ratio0.75、RandomFlip与PhotoMetricDistortion评测指标为IoUMetric的mIoU调度基座是 configs/base/schedules/schedule_80k.py默认 80000 迭代训练每 8000 迭代验证并保存 checkpointSegDataPreProcessor采用 ImageNet 统计的均值/标准差bgr_to_rgbTrue说明输入按 BGR 读取后转 RGBseg_pad_val255表示标注 padding 值使用 ignore index。模型主体ResNet-50-D8 IterativeDecodeHeadnorm_cfg dict(typeSyncBN, requires_gradTrue) num_stages 3 conv_kernel_size 1 model dict( typeEncoderDecoder, data_preprocessordata_preprocessor, pretrainedopen-mmlab://resnet50_v1c, backbonedict( typeResNetV1c, depth50, num_stages4, out_indices(0, 1, 2, 3), dilations(1, 1, 2, 4), strides(1, 2, 1, 1), norm_cfgnorm_cfg, norm_evalFalse, stylepytorch, contract_dilationTrue), decode_headdict( typeIterativeDecodeHead, num_stagesnum_stages, kernel_update_head[ dict( typeKernelUpdateHead, num_classes150, num_ffn_fcs2, num_heads8, num_mask_fcs1, feedforward_channels2048, in_channels512, out_channels512, dropout0.0, conv_kernel_sizeconv_kernel_size, ffn_act_cfgdict(typeReLU, inplaceTrue), with_ffnTrue, feat_transform_cfgdict(conv_cfgdict(typeConv2d), act_cfgNone), kernel_updator_cfgdict( typeKernelUpdator, in_channels256, feat_channels256, out_channels256, act_cfgdict(typeReLU, inplaceTrue), norm_cfgdict(typeLN))) for _ in range(num_stages) ], kernel_generate_headdict( typeASPPHead, in_channels2048, in_index3, channels512, dilations(1, 12, 24, 36), dropout_ratio0.1, num_classes150, norm_cfgnorm_cfg, align_cornersFalse, loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, loss_weight1.0))), auxiliary_headdict( typeFCNHead, in_channels1024, in_index2, channels256, num_convs1, concat_inputFalse, dropout_ratio0.1, num_classes150, norm_cfgnorm_cfg, align_cornersFalse, loss_decodedict( typeCrossEntropyLoss, use_sigmoidFalse, loss_weight0.4)), train_cfgdict(), test_cfgdict(modewhole))关键参数含义参数取值说明num_stages3迭代内核更新的阶段数必须与kernel_update_head列表长度一致源码中assert num_stages len(kernel_update_head)conv_kernel_size1动态内核的卷积核尺寸1 即退化为 1×1 卷积也是配置的默认选择num_ffn_fcs2Kernel Update Head 中 FFN 的全连接层数num_heads8内核自注意力的注意力头数num_mask_fcs1掩码预测前的全连接层组数feedforward_channels2048FFN 的隐层宽度in_channels/out_channels512内核特征维度与生成头输出通道对齐kernel_updator_cfgKernelUpdator内核更新器内部in/feat/out_channels256激活 ReLU、归一化 LayerNormkernel_generate_headASPPHead初始内核与掩码的生成头dilations(1,12,24,36) 的多尺度空洞卷积池auxiliary_headFCNHead辅助头loss_weight0.4 辅助监督中低层特征注意 backbone 的dilations(1, 1, 2, 4)与strides(1, 2, 1, 1)组合即 D8 空洞策略stage3/4 空洞率为 2/4维持较高分辨率的特征输出out_indices(0,1,2,3)输出全部四个阶段供生成头与辅助头使用。优化器与学习率策略K-Net 特有optim_wrapper dict( _delete_True, typeOptimWrapper, optimizerdict(typeAdamW, lr0.0001, weight_decay0.0005), clip_graddict(max_norm1, norm_type2)) param_scheduler [ dict(typeLinearLR, start_factor0.001, by_epochFalse, begin0, end1000), dict(typeMultiStepLR, begin1000, end80000, milestones[60000, 72000], by_epochFalse) ]与 schedule_80k.py 基座的 SGD PolyLR 不同K-Net 配置通过_delete_True彻底替换为AdamWlr1e-4weight_decay5e-4并开启梯度裁剪max_norm1学习率采用1000 步线性 warmupstart_factor0.001 MultiStepLR 在 60000/72000 步衰减的组合策略。数据加载批量设置# In K-Net implementation we use batch size 2 per GPU as default train_dataloader dict(batch_size2, num_workers2) val_dataloader dict(batch_size1, num_workers4) test_dataloader val_dataloader配置注释明确说明 K-Net 默认每 GPU 批量 2。结合 README 的说明所有实验在 8×V100 32G 上、每 GPU 2 个采样器8 GPU × 2 16 的总批量与 metafile.yaml 中记录的Batch Size: 16完全一致。四种内核生成头的组合与 Swin 变体K-Net 的“即插即用”体现在kernel_generate_head可以替换为任意现有分割头从而复用其 FPN 式多尺度特征与成熟的监督设计。仓库 configs/knet 提供了四种 ResNet-50 组合与两种 Swin 变体配置生成头backbone 细节knet-s3_r50-d8_fcn_8xb2-adamw-80k_ade20k-512x512.pyFCNHeadin_channels2048, num_convs2, concat_inputTrueR-50-D8knet-s3_r50-d8_pspnet_8xb2-adamw-80k_ade20k-512x512.pyPSPHeadpool_scales(1,2,3,6)R-50-D8knet-s3_r50-d8_deeplabv3_8xb2-adamw-80k_ade20k-512x512.pyASPPHeaddilations(1,12,24,36)R-50-D8knet-s3_r50-d8_upernet_8xb2-adamw-80k_ade20k-512x512.pyUPerHeadin_channels[256,512,1024,2048], pool_scales(1,2,3,6)R-50-D8strides(1,2,2,2) 非空洞knet-s3_swin-t_upernet_8xb2-adamw-80k_ade20k-512x512.pyUPerHeadin_channels[96,192,384,768]Swin-Tdrop_path_rate0.3knet-s3_swin-l_upernet_8xb2-adamw-80k_ade20k-512x512.pyUPerHeadin_channels[192,384,768,1536]Swin-Ldrop_path_rate0.4knet-s3_swin-l_upernet_8xb2-adamw-80k_ade20k-640x640.pyUPerHeadin_channels[192,384,768,1536]Swin-L640×640 训练尺度Swin 变体在继承前一级配置的基础上通过_delete_True替换 backbone 与优化器学习率改为 0.00006遵循 Swin Transformer 官方实现并对absolute_pos_embed、relative_position_bias_table、norm设置decay_mult0的权重衰减豁免。knet-s3_swin-l_upernet_8xb2-adamw-80k_ade20k-640x640.py 还覆写了完整训练/测试管线RandomResize的 scale 为 (2048, 640)RandomCrop的 crop_size 为 (640, 640)。ADE20K 实验结果以下结果表完整引自 configs/knet/README.md各实验均基于 8×V100 32G每 GPU 批量 2共 16 的批量大小MethodBackboneCrop SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)configKNet FCNR-50-D8512x512800007.0119.24V10043.6045.12configKNet PSPNetR-50-D8512x512800006.9820.04V10044.1845.58configKNet DeepLabV3R-50-D8512x512800007.4212.10V10045.0646.11configKNet UperNetR-50-D8512x512800007.3417.11V10043.4544.07configKNet UperNetSwin-T512x512800007.5715.56V10045.8446.27configKNet UperNetSwin-L512x5128000013.58.29V10052.0553.24configKNet UperNetSwin-L640x6408000013.548.29V10052.2153.34config数据解读生成头之间的差异在同一 R-50-D8 骨干下ASPP 生成头DeepLabV3 组合取得最高 mIoU 45.06说明多尺度空洞上下文对初始内核质量更有帮助而 FCN/PSPNet/UperNet 组合保持在 43.4–44.2 区间骨干的收益Swin-T 将 UperNet 组合提升至 45.84Swin-L512×512达 52.05640×640 训练尺度进一步微升至 52.21msflip 评测 53.34同时显存从约 7.5GB 升至 13.5GB多尺度 翻转测试msflip 相比单尺度普遍提升约 1.1–1.5 个点ADE20K 基座配置中自带的tta_pipeline6 个尺度 × 水平翻转即用于该评测。上述每个模型的预训练权重与训练日志均登记在 configs/knet/metafile.yaml 中含 43.6/45.12、44.18/45.58、45.06/46.11、43.45/44.07、45.84/46.27、52.05/53.24、52.21/53.34 等指标的官方记录框架信息标注为 PyTorch许可证为 Apache License 2.0同时汇总于仓库根目录的 model-index.yml可用于模型下载与自动评测。训练与评测实操使用仓库自带的入口脚本即可复现 K-Net。单机多卡训练8 GPU 与 README 记录一致bash tools/dist_train.sh configs/knet/knet-s3_r50-d8_deeplabv3_8xb2-adamw-80k_ade20k-512x512.py 8单卡训练可直接调用 tools/train.pypython tools/train.py configs/knet/knet-s3_r50-d8_deeplabv3_8xb2-adamw-80k_ade20k-512x512.py评测使用 tools/test.py加载metafile.yaml中记录的对应 checkpoint 即可得到 mIoUpython tools/test.py configs/knet/knet-s3_r50-d8_deeplabv3_8xb2-adamw-80k_ade20k-512x512.py /path/to/knet_checkpoint.pth --eval mIoU实操提醒以当前仓库实际内容为准ADE20K 数据集需按 configs/base/datasets/ade20k.py 的约定放置于data/ade/ADEChallengeData2016并组织为images/training、annotations/training、images/validation、annotations/validation目录结构所有 K-Net 配置默认test_cfgdict(modewhole)即整图推理如需滑动窗口推理可参照其他算法配置调整若显存受限可降低train_dataloader的batch_size但需注意 README 中的指标是在 8×V10032G、每 GPU 批量 2 的条件下取得的。引用信息如果 K-Net 对你的工作有所帮助请按 configs/knet/README.md 中提供的 BibTeX 引用该论文inproceedings{zhang2021knet, title{{K-Net: Towards} Unified Image Segmentation}, author{Wenwei Zhang and Jiangmiao Pang and Kai Chen and Chen Change Loy}, year{2021}, booktitle{NeurIPS}, }小结K-Net 以“一组随图像内容动态更新的可学习内核”统一了语义/实例/全景分割的任务表述而 MMSegmentation 通过IterativeDecodeHeadKernelUpdateHeadKernelUpdator三件套在 knet_head.py 中给出了工业级实现并提供了从 FCN 到 DeepLabV3、从 ResNet 到 Swin 的多组即用配置。理解本文梳理的源码调用链与配置参数你便可以在 configs/knet 的基础上自由组合骨干与生成头将 K-Net 迁移到自己的语义分割任务中。【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进