MMSegmentation 中的 BEiT 骨干网络：从权重转换到 ADE20K 语义分割实战

MMSegmentation 中的 BEiT 骨干网络从权重转换到 ADE20K 语义分割实战【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentationBEiTBidirectional Encoder representation from Image Transformers是微软提出的基于掩码图像建模的自监督视觉 Transformer 预训练方法。本文以 configs/beit/README.md 为主线结合 MMSegmentation 仓库中的源码与配置文件完整讲解如何在当前仓库中使用 BEiT 作为骨干网络进行语义分割包括预训练权重键名转换、UPerNet 训练配置解读、单尺度与多尺度滑动窗口推理以及 ADE20K 上的公开基准结果。BEiT 方法简介BEiT 论文发表于 ICLR 2022arXiv:2106.08254核心思想是将自然语言处理中 BERT 的预训练范式迁移到视觉领域。其预训练流程包含两个视图图像块image patches将图像切分为例如 16×16 像素的块视觉词元visual tokens通过离散 tokenizer 将原始图像分词为离散 token。预训练时随机掩码一部分图像块输入骨干 Transformer训练目标是从被破坏的图像块中恢复原始视觉 token。预训练完成后直接在编码器上追加任务层对下游任务进行微调。论文报告的分类结果表明该方法的竞争力base 规模的 BEiT 在 ImageNet-1K 上达到 83.2% top-1 准确率明显优于相同设置下从零训练的 DeiT81.8%large 规模仅用 ImageNet-1K 即达到 86.3%甚至超过在 ImageNet-22K 上监督预训练的 ViT-L85.2%。MMSegmentation 当前仓库在configs/beit/目录下提供了 BEiT-Base 与 BEiT-Large 搭配 UPerNet 的完整语义分割方案覆盖 ADE20K 数据集的训练与测试配置。MMSegmentation 中的 BEiT 骨干实现骨干网络源码结构BEiT 骨干网络实现在 mmseg/models/backbones/beit.py共 554 行包含以下核心模块BEiTAttention基于窗口的多头自注意力W-MSA并带有相对位置偏置。其bias参数支持qv_bias仅给 q、v 加可学习偏置、Trueq、k、v 都加与False不加三种模式相对位置编码_init_rel_pos_embedding中构建(2*Wh-1)*(2*Ww-1)3大小的相对位置偏置表额外 3 项分别对应 cls-to-token、token-to-cls、cls-to-cls 的相对距离TransformerEncoderLayer 复用文件头部from .vit import TransformerEncoderLayer as VisionTransformerEncoderLayer说明其编码器层复用了 mmseg/models/backbones/vit.py 中的实现。从源码结构可以推断BEiT 在分割场景中与 ViT 共享大量 Transformer 基础组件差异主要集中在带相对位置偏置的窗口注意力机制上——这正是 BEiT 与普通 ViT 的关键区别。骨干注册与配置入口BEiT 骨干以typeBEiT形式注册在MODELSregistry 中通过 configs/base/models/upernet_beit.py 提供默认配置backbonedict( typeBEiT, img_size(640, 640), patch_size16, in_channels3, embed_dims768, num_layers12, num_heads12, mlp_ratio4, out_indices(3, 5, 7, 11), qv_biasTrue, attn_drop_rate0.0, drop_path_rate0.1, norm_cfgdict(typeLN, eps1e-6), act_cfgdict(typeGELU), norm_evalFalse, init_values0.1),各参数含义参数说明img_size输入图像尺寸语义分割中固定为 (640, 640)patch_size图像块大小BEiT 使用 16×16embed_dims特征嵌入维度Base 为 768Large 为 1024num_layers/num_headsTransformer 层数与注意力头数Base 为 12/12Large 为 24/16mlp_ratioFFN 隐藏层放大比例此处为 4out_indices输出特征层索引取第 3、5、7、11或 Large 的第 7、11、15、23层作为多尺度特征qv_bias仅给 q、v 加可学习偏置对应论文实现init_valuesLayerScale 初始值UPerNet 分割头通过neckdict(typeFeature2Pyramid, embed_dim768, rescales[4, 2, 1, 0.5])将四层同维特征BEiT 各层输出维度相同均为embed_dims通过不同缩放比例重建为金字塔特征再送入UPerHead解码。预训练权重转换beit2mmseg.py官方 BEiT 预训练权重的键名key与 MMSegmentation 风格不一致直接加载会报错。仓库在 tools/model_converters/beit2mmseg.py 提供了键名转换脚本其核心逻辑convert_beit完成三类映射官方键名模式转换后键名patch_embed.proj.*patch_embed.projection.*blocks.*.norm.*layers.*.ln.*blocks.*.mlp.fc1.*/blocks.*.mlp.fc2.*layers.*.ffn.layers.0.0.*/layers.*.ffn.layers.1.*blocks.*其余子键layers.*脚本使用mmengine.runner.CheckpointLoader.load_checkpoint加载权重自动兼容state_dict、model或裸权重三种 checkpoint 格式转换后经torch.save保存到目标路径。使用方法python tools/model_converters/beit2mmseg.py ${PRETRAIN_PATH} ${STORE_PATH}PRETRAIN_PATH为官方预训练权重路径或 URLSTORE_PATH为转换后权重的保存路径。README 给出的实际示例python tools/model_converters/beit2mmseg.py https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22k.pth pretrain/beit_base_patch16_224_pt22k_ft22k.pth当前仓库默认设置下使用的两个预训练权重及其来源对应关系如下转换后权重原始官方权重BEiT_base.pthBEiT_basepatch16, 224×224, ImageNet-22K 预训练微调BEiT_large.pthBEiT_largepatch16, 224×224, ImageNet-22K 预训练微调转换完成后将权重放置在配置中pretrained字段指定的路径默认pretrain/目录即可被训练/测试流程加载。训练配置详解BEiT-Base UPerNetconfigs/beit/beit-base_upernet_8xb2-160k_ade20k-640x640.py 是 Base 模型的训练配置继承自四个基础配置../_base_/models/upernet_beit.pyBEiTUPerNet 模型定义../_base_/datasets/ade20k_640x640.pyADE20K 数据集配置../_base_/default_runtime.py默认运行环境../_base_/schedules/schedule_160k.py160k 迭代训练计划。关键训练设置optim_wrapper dict( _delete_True, typeOptimWrapper, optimizerdict( typeAdamW, lr3e-5, betas(0.9, 0.999), weight_decay0.05), constructorLayerDecayOptimizerConstructor, paramwise_cfgdict(num_layers12, layer_decay_rate0.9)) param_scheduler [ dict( typeLinearLR, start_factor1e-6, by_epochFalse, begin0, end1500), dict( typePolyLR, power1.0, begin1500, end160000, eta_min0.0, by_epochFalse, ) ]要点解读采用AdamW 优化器学习率 3e-5权重衰减 0.05_delete_True表示覆盖基配置中的默认优化器LayerDecayOptimizerConstructor分层衰减layer_decay_rate0.9num_layers12浅层参数学习率更低、深层更高这是 Transformer 微调中抑制浅层剧烈更新的常用策略两阶段学习率调度前 1500 次迭代用LinearLR起始因子 1e-6做 warmup之后切换为power1.0的PolyLR直到 160k 迭代结束数据并行默认 8 张 GPU、每卡 2 张图train_dataloader dict(batch_size2)即有效 batch size 为 16验证与测试每卡 1 张。BEiT-Large UPerNetAMP 训练configs/beit/beit-large_upernet_8xb1-amp-160k_ade20k-640x640.py 是 Large 模型配置与 Base 版本的关键差异配置项BEiT-BaseBEiT-Largeembed_dims7681024num_layers1224num_heads1216drop_path_rate0.10.2out_indices(3, 5, 7, 11)(7, 11, 15, 23)layer_decay_rate0.90.95学习率3e-52e-5优化器包装OptimWrapperAmpOptimWrapper每卡 batch size21accumulative_counts2梯度累积训练计划schedule_160kschedule_320kLarge 配置通过AmpOptimWrapper启用自动混合精度AMP训练配合accumulative_counts2梯度累积在每卡仅 1 张图8 卡有效 batch size 为 8的条件下模拟更大的 batch。由于模型更深24 层drop_path_rate提高到 0.2 以增强正则化分层衰减率提高到 0.95。注意Large 配置中学习率调度器的end为 160000而继承的计划为 320k 迭代——从配置结构看这属于该文件的既有设定使用时请以实际迭代计划为准README 结果表中 Large 记为 320000 迭代。ADE20K 数据与预处理configs/base/datasets/ade20k_640x640.py 定义训练数据流随机缩放RandomResize以(2560, 640)为基准尺度、ratio_range(0.5, 2.0)随机缩放并保持比例随机裁剪RandomCrop裁剪到 640×640cat_max_ratio0.75限制裁剪区域内单一类别占比不超过 75%避免裁剪到大面积同质区域数据增强RandomFlip概率 0.5与PhotoMetricDistortion光度扰动类别处理LoadAnnotations使用reduce_zero_labelTrueADE20K 索引 0 的背景类被折叠因此模型输出num_classes150。训练数据前缀为images/training与annotations/training数据根目录默认data/ade/ADEChallengeData2016需要用户预先按该结构准备数据。测试与推理单尺度测试README 提供的单尺度验证命令权重文件名为当前仓库发布版对应名称命令中的旧配置文件名在新版本中已更名为上文所述配置sh tools/dist_test.sh \ configs/beit/beit-large_upernet_8xb1-amp-160k_ade20k-640x640.py \ upernet_beit-large_fp16_8x1_640x640_160k_ade20k-8fc0dd5d.pth $GPUS --eval mIoU$GPUS为使用的 GPU 数量--eval mIoU指定评估指标为 mIoU。多尺度滑动窗口测试由于 BEiT 的相对位置编码要求输入的长和宽相等MMSegmentation 采用**滑动窗口slide**策略进行多尺度推理将最短边固定为 640min_size640因此多尺度推理需要单独执行_ms配置而不是使用通用的--aug-test参数sh tools/dist_test.sh \ configs/beit/beit-large_upernet_8xb1-amp-160k_ade20k-640x640_ms.py \ upernet_beit-large_fp16_8x1_640x640_160k_ade20k-8fc0dd5d.pth $GPUS --eval mIoU对应的多尺度配置 configs/beit/beit-large_upernet_8xb1-amp-160k_ade20k-640x640_ms.py 通过覆盖test_pipeline实现test_pipeline [ dict(typeLoadImageFromFile), # img_ratios is [0.5, 0.75, 1.0, 1.25, 1.5, 1.75] # original image scale is (2560, 640) dict(typeResize, scale(2560, 640), keep_ratioTrue), dict(typeLoadAnnotations, reduce_zero_labelTrue), dict(typePackSegInputs), ]推理流程为将原图按(2560, 640)缩放保持比例再以test_cfgdict(modeslide, crop_size(640, 640), stride(426, 426))滑动窗口切块推理并拼接。其中 stride 取 640 的 2/3426保证窗口间有重叠缓解边界伪影。Base 与 Large 配置的test_cfg均为该滑动窗口设置。ADE20K 基准结果README 记录了 UPerNet 搭配两种 BEiT 骨干在 ADE20K 验证集上的结果MethodBackboneCrop SizepretrainBatch SizeLr schdMem (GB)Inf time (fps)DevicemIoUmIoU(msflip)UPerNetBEiT-B640x640ImageNet-22K1616000015.882.00V10053.0853.84UPerNetBEiT-L640x640ImageNet-22K832000022.640.96V10056.3356.84对应训练配置为 beit-base_upernet_8xb2-160k_ade20k-640x640.py 与 beit-large_upernet_8xb1-amp-160k_ade20k-640x640.py。其中msflip列对应多尺度水平翻转融合的结果即通过_ms配置与滑动窗口推理获得。模型权重与训练日志的下载地址以仓库 configs/beit/metafile.yaml 中登记的信息为准。总结BEiT 在 MMSegmentation 中的落地路径清晰完整通过 beit2mmseg.py 完成官方权重键名转换复用 beit.py 中带相对位置偏置的窗口注意力骨干以 UPerNet 为解码头在 ADE20K 上达到 53.08BEiT-B与 56.33BEiT-L的 mIoU。实操中的两个关键经验是相对位置编码要求方形输入多尺度推理必须走滑动窗口配置Transformer 微调建议使用 LayerDecay 分层衰减与较小的 AdamW 学习率。引用若在研究中引用 BEiT请使用以下 BibTeXinproceedings{beit, title{{BEiT}: {BERT} Pre-Training of Image Transformers}, author{Hangbo Bao and Li Dong and Songhao Piao and Furu Wei}, booktitle{International Conference on Learning Representations}, year{2022}, url{https://openreview.net/forum?idp-BhZSz59o4} }【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考