BEiT v2 预训练完全指南：在 ImageNet-1k 上从零训练掩码图像建模（MIM）模型

BEiT v2 预训练完全指南在 ImageNet-1k 上从零训练掩码图像建模MIM模型【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilmBEiT v2 是微软 unilm 仓库中基于向量量化视觉 TokenizerVQ-KD的掩码图像建模Masked Image Modeling, MIM预训练方案本指南以beit2/PRETRAINING.md为核心结合仓库源码完整讲解如何从零在 ImageNet-1k 上预训练 BEiT v2 base/large 模型包括分布式启动方式、全部关键超参数与默认值、cls-token 预训练结构early/head layers、VQ-KD Tokenizer 的加载与使用以及官方发布的预训练权重Model Zoo。读完本文你将能直接复制并运行一条可用的 BEiT v2 预训练命令并理解每个参数在源码中的具体作用。背景BEiT v2 预训练的两个核心组件BEiT v2 预训练延续了 BEiT v1 的思路原文档明确说明 We follow the settings proposed in BEiT v1但核心区别在于预训练的目标不再是对原始像素做回归而是用向量量化的视觉 Tokenizer 把图像转成离散 token再让 ViT 去还原被掩码的 token。从源码结构看预训练阶段涉及两个模型BEiT v2 主干backbone定义在 beit2/modeling_pretrain.py是一个带mask_token的 Vision TransformerVQ-KD 视觉 Tokenizer定义在 beit2/modeling_vqkd.py编码器 量化器 解码器结构负责把图像压缩成 8192 词表codebook中的离散 token 序列。预训练入口脚本是 beit2/run_beitv2_pretraining.py单轮训练的核心流程实现于 beit2/engine_for_pretraining.py。预训练环境准备与依赖根据 beit2/README.md需要先安装 beit2/requirements.txt 中锁定的依赖版本PyTorch 1.7.1torchvision 0.8.2timm 0.4.12deepspeed 0.4.0可选用于训练加速einops、tensorboardX、blobfile、requests 等辅助库安装方式cd beit2 pip install -r requirements.txt若需要使用混合精度训练还需额外安装 NVIDIA apex--cpp_ext --cuda_ext选项。注意以上版本要求针对仓库当前代码环境若使用更新的 PyTorch/timm 版本需要自行验证兼容性。核心一在 ImageNet-1k 上预训练 BEiT v2 Base 模型原文档给出 Base 模型在 1 台 DGX-216 块 V100-32GB上的完整命令这是最经典的参考配置python -m torch.distributed.launch --nproc_per_node16 run_beitv2_pretraining.py \ --data_set image_folder \ --data_path /path/to/imagenet-1k/train \ --output_dir /path/to/save/your_model \ --log_dir /path/to/save/your_model \ --model beit_base_patch16_224_8k_vocab_cls_pt \ --shared_lm_head True \ --early_layers 9 \ --head_layers 2 \ --num_mask_patches 75 \ --second_input_size 224 \ --second_interpolation bicubic \ --min_crop_scale 0.2 \ --tokenizer_model vqkd_encoder_base_decoder_3x768x12_clip \ --tokenizer_weight https://github.com/addf400/files/releases/download/BEiT-v2/vqkd_encoder_base_decoder_3x768x12_clip-d5036aa7.pth \ --batch_size 128 \ --lr 1.5e-3 \ --warmup_epochs 10 \ --clip_grad 3.0 \ --drop_path 0.1 \ --layer_scale_init_value 0.1 \ --imagenet_default_mean_and_std \ --opt_betas 0.9 0.999 \ --opt_eps 1e-8 \ --epochs 1600 \ --save_ckpt_freq 20关键参数逐项说明结合源码--model主干模型名。beit_base_patch16_224_8k_vocab_cls_pt表示带 cls-token 预训练的 base 模型对应源码中的VisionTransformerForMaskedImageModelingCLS见 beit2/modeling_pretrain.pybeit_base_patch16_224_8k_vocab表示不带 cls-token 预训练的版本即 BEiT v1 base 模型对应VisionTransformerForMaskedImageModeling。--batch_size每张 GPU 的 batch size。有效 batch size GPU 数 ×--batch_size因此上例有效 batch size 为128 × 16 2048。--tokenizer_model推荐使用vqkd_encoder_base_decoder_3x768x12_clip编码器 12 层、解码器 3 层、隐藏维度 768。--tokenizer_weightTokenzier 权重的下载路径或本地路径。从源码看run_beitv2_pretraining.py中get_visual_tokenizer()会通过create_model(..., pretrainedTrue, pretrained_weightargs.tokenizer_weight, as_tokenzerTrue)加载它并且支持以https://开头的远程地址内部走torch.hub.load_state_dict_from_url加载时会剔除loss、teacher、scaling前缀参数只保留 encoder/decoder/quantize 权重见 beit2/modeling_vqkd.py。--epochs短日程用 300长日程用 1600。--opt_betas300 epoch 用0.981600 epoch 用0.999上例属于长日程故为 0.9/0.999。--drop_path300 epoch 用0.01600 epoch 用0.1。--early_layers/--head_layers/--shared_lm_headcls-token 预训练专属结构。从VisionTransformerForMaskedImageModelingCLS的实现看主干前early_layersbase 为 9large 为 21层正常前向之后把cls_token与第early_layers层的 patch 特征拼接再过head_layers默认为 2层condenser head最终两条支路共享同一个lm_head预测被掩码 tokenshared_lm_headTrue时两条支路共用一个分类头为 False 时额外新建cls_pt_norm与cls_pt_lm_headbeit2/modeling_pretrain.py。--layer_scale_init_valuebase 用 0.1large 用 1e-5设 0 表示禁用 layer scale源码默认 0.1。--num_mask_patches被掩码的 patch 数默认 75配合--min_mask_patches_per_block默认 16与--max_mask_patches_per_block默认 None即不超过剩余掩码数控制块状掩码生成掩码算法实现于 beit2/masking_generator.py。--second_input_size/--second_interpolation喂给 VQ-KD Tokenzier 的第二个视角的分辨率与插值方式。数据增强中通过RandomResizedCropAndInterpolationWithTwoPic同时产出两份裁剪结果一份经归一化后进主干一份不进归一化直接进 Tokenzier见 beit2/datasets.py。--min_crop_scale随机裁剪的最小面积比例预训练建议 0.2源码默认 0.08微调/常规分类默认更小。--imagenet_default_mean_and_std使用 ImageNet 默认归一化均值/方差不开启时使用 Inception 风格的 mean/std见 beit2/datasets.py。--clip_grad梯度裁剪范数 3.0。--save_ckpt_freq每 N 个 epoch 保存一次 checkpoint默认 20。预训练单轮训练流程源码级在 beit2/engine_for_pretraining.py 的train_one_epoch中每个 iteration 的执行顺序是从 dataloader 取(samples, images, bool_masked_pos)samples是归一化后的图像给主干images是未归一化图像给 Tokenziervqkd.get_codebook_indices(images)在torch.no_grad()与 AMP 下把图像编码成离散 token作为回归标签input_ids主干对掩码后的 patch 前向输出经lm_head的预测 logits使用nn.CrossEntropyLoss计算损失若启用了 cls-token 预训练outputs是一个 list最终 loss 为两条支路之和loss_1 loss_2并分别记录mlm_acc_1/mlm_acc_2即每条支路对掩码 token 的预测准确率。预训练脚本主流程beit2/run_beitv2_pretraining.py还包含分布式初始化utils.init_distributed_mode、step 级余弦学习率与 weight decay 调度utils.cosine_scheduler、NativeScaler带梯度范数的 AMP loss scaler以及utils.auto_load_model的断点续训与--resume支持。核心二预训练 BEiT v2 Large 模型4 机分布式Large 模型需要 4 台 DGX-24×16 块 V100-32GB使用torch.distributed.launch的多节点模式python -m torch.distributed.launch --nnodes 4 --node_rank {0, 1, 2, 3} --nproc_per_node16 run_beitv2_pretraining.py \ --data_set image_folder \ --data_path /path/to/imagenet-1k/train \ --output_dir /path/to/save/your_model \ --log_dir /path/to/save/your_model \ --model beit_large_patch16_224_8k_vocab_cls_pt \ --shared_lm_head True \ --early_layers 21 \ --head_layers 2 \ --num_mask_patches 75 \ --second_input_size 224 \ --second_interpolation bicubic \ --min_crop_scale 0.2 \ --tokenizer_model vqkd_encoder_base_decoder_3x768x12_clip \ --tokenizer_weight https://github.com/addf400/files/releases/download/BEiT-v2/vqkd_encoder_base_decoder_3x768x12_clip-d5036aa7.pth \ --batch_size 32 \ --lr 1.5e-3 \ --warmup_epochs 10 \ --clip_grad 3.0 \ --drop_path 0.1 \ --layer_scale_init_value 1e-5 \ --imagenet_default_mean_and_std \ --opt_betas 0.9 0.999 \ --opt_eps 1e-8 \ --epochs 1600 \ --save_ckpt_freq 20要点与 Base 的区别有效 batch size32 × (4 × 16) 2048与 Base 示例保持一致--modelbeit_large_patch16_224_8k_vocab_cls_pt为 large cls-token 预训练beit_large_patch16_224_8k_vocab为不带 cls-token 的 large即 BEiT v1 large--early_layerslarge 用 21源码默认注释 default 9 for base and 21 for large--layer_scale_init_valuelarge 用1e-5源码注释 0.1 for base, 1e-5 for large--node_rank在 4 个节点上分别指定 0/1/2/3其余参数保持一致。模型结构上large 在源码中对应embed_dim1024, depth24, num_heads16beit2/modeling_pretrain.py而 base 为embed_dim768, depth12, num_heads12。预训练权重Model Zoo原文档提供 4 个在 ImageNet-1k 上自监督预训练完成的官方权重可直接用于后续微调如 beit2/get_started_for_image_classification.md 中的分类微调--finetune参数model namepretraining epochs说明beit_base_patch16_224_8k_vocab_cls_pt300base cls-token短日程 300 epochbeit_base_patch16_224_8k_vocab_cls_pt1600base cls-token长日程 1600 epochbeit_large_patch16_224_8k_vocab_cls_pt300large cls-token短日程 300 epochbeit_large_patch16_224_8k_vocab_cls_pt1600large cls-token长日程 1600 epoch结合仓库 README还可以进一步获取ImageNet-1k 预训练 ImageNet-21k 中间微调的推荐权重beitv2_base_patch16_224_pt1k_ft21k等以及微调后的分类/分割权重。配套核心组件VQ-KD Tokenizer 与图像编码预训练命令中的--tokenizer_model vqkd_encoder_base_decoder_3x768x12_clip指向官方在 ImageNet-1k 上训练好的 VQ-KD Tokenzier。相关细节见 beit2/TOKENIZER.md结构vqkd_encoder_base_decoder_3x768x12_clip表示编码器 12 层、解码器 3 层、hidden768、FFN 4x、head12、patch16×16Teacher 为 CLIP ViT-B/16vqkd_encoder_base_decoder_1x768x12_dino则使用 DINO ViT-B/16 作为 Teacher解码器输出维度变为 768原理VQ-KD 的目标是从 Teacher 重建语义知识而非原始像素从而构造紧凑的语义码本codebook8192 词、32 维源码中量化器为带 EMA 更新与 k-means 初始化的NormEMAVectorQuantizerbeit2/norm_ema_quantizer.py重建损失为余弦相似度损失使用方式预训练脚本直接调用vqkd.get_codebook_indices(images)得到 token 序列beit2/modeling_vqkd.py如果想独立把图像压缩成离散 token可以运行 beit2/test_get_code.pypython test_get_code.py。如果需要自行训练 VQ-KD Tokenzier原仓库也提供了run_vqkd_training.py的完整命令与超参数--codebook_n_emd 8192、--codebook_emd_dim 32、--quantize_kmeans_init、--rec_loss_type cosine、Teacher 为 CLIP 等详见 beit2/TOKENIZER.md。常见问题与调参速查显存不足时怎么降有效 batch size优先降低--batch_size每卡保持GPU 数 × batch_size 2048的经验值若单卡显存有限可配合 deepspeed--enable_deepspeed与梯度累积。短日程与长日程如何切换300 epoch 短日程建议--opt_betas 0.98、--drop_path 0.01600 epoch 长日程建议--opt_betas 0.999、--drop_path 0.1其余结构参数保持一致。为什么需要--second_input_size因为同一张图要同时喂给主干归一化与 Tokenzier不归一化两份输入的裁剪尺寸可以不同两者插值方式也可分别用--train_interpolation与--second_interpolation控制。--model结尾的cls_pt是什么意思表示使用带 cls-token 的 condenser 预训练结构BEiT v2 新特性不带cls_pt的模型即 BEiT v1 结构。选用时务必同时设置对应的--early_layersbase9 / large21与--shared_lm_head。如何断点续训使用--resume指定 checkpoint或开启--auto_resume源码默认开启会自动从output_dir恢复。预训练之后预训练产物--output_dir下按--save_ckpt_freq周期保存的 checkpoint可直接用于下游任务微调图像分类微调流程与完整参数见 beit2/get_started_for_image_classification.md其中--finetune可填本仓库 Model Zoo 的预训练权重或自己训练得到的权重语义分割微调见 beit2/semantic_segmentation/README.md。从 README 的微调结果看base1600 epoch 预训练在 ImageNet-1k 上可达到 85.5% top-1 准确率large 可达 87.3%若经过 ImageNet-21k 中间微调则进一步提升——这些数据均来自仓库 beit2/README.md 的记录可作为预训练质量的对标参考。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考