Detectron2 中重思 BatchNorm 的 Batch 语义Rethinking-BatchNorm 项目配置与源码深度解析【免费下载链接】detectron2Detectron2 is a platform for object detection, segmentation and other visual recognition tasks.项目地址: https://gitcode.com/GitHub_Trending/de/detectron2导读BatchNorm 的效果高度依赖 batch 的统计口径——这里的 batch 到底应该指单卡上的 mini-batch、跨卡聚合的全局 batch还是某个特征域feature domain自身的统计本文以 Detectron2 仓库中 Rethinking-BatchNorm 项目为骨架逐条解析其 6 份可复现论文《Rethinking Batch in BatchNorm》实验的 LazyConfig 配置文件与评测脚本并结合detectron2/layers/batch_norm.py、detectron2/modeling/meta_arch/retinanet.py等源码说明底层机制。读完你将掌握在 Mask R-CNN / RetinaNet 的 head 中切换 BN 语义单卡 BN、batch 统计量、跨卡 shuffle、SyncBN、共享 BN、域特定 BN的完整配置方法以及如何用域特定统计量脚本复现论文 Table 5 的高精度结果。一、背景BatchNorm 的 Batch 到底指什么BatchNorm 在训练时用当前 batch 的均值/方差归一化激活并在推理时切换为滑动平均得到的全局统计量。但在分布式训练、多尺度特征图等场景下当前 batch 的边界是模糊的单卡 mini-batchbatch 只包含当前 GPU 上的样本batch 较小时统计量噪声大跨卡全局 batchSyncBatchNorm 把多卡统计量聚合后再归一化等价于扩大了 batch特征域domainbatch同一层可能被多个输入域复用如 RetinaNet 的 5 个特征金字塔层共享同一个 head每个域应维护自己的测试期统计量。论文 Rethinking Batch in BatchNorm 系统研究了这些问题而本仓库 projects/Rethinking-BatchNorm 提供了一套 LazyConfig 实验配置用于在 Detectron2 上复现论文中的 Mask R-CNNTable 3、Figure 9、Table 6与 RetinaNetTable 5检测实验。二、项目结构与快速上手2.1 目录构成projects/Rethinking-BatchNorm/ ├── configs/ │ ├── mask_rcnn_BNhead.py # Mask R-CNNhead 中使用 BatchNorm │ ├── mask_rcnn_BNhead_batch_stats.py # 推理时改用 batch 统计量的 BN │ ├── mask_rcnn_BNhead_shuffle.py # 跨 GPU 打乱 head 输入 │ ├── mask_rcnn_SyncBNhead.py # head 中使用 SyncBN │ ├── retinanet_SyncBNhead.py # RetinaNet head 使用 SyncBN │ └── retinanet_SyncBNhead_SharedTraining.py # 5 个特征层共享归一化统计 ├── retinanet-eval-domain-specific.py # 重算域特定统计量的评测脚本 └── README.md2.2 训练命令所有配置都可以直接通过 Detectron2 的 LazyConfig 训练入口启动按 README 中命令在projects/Rethinking-BatchNorm/目录下执行../../tools/lazyconfig_train_net.py --config-file configs/X.py --num-gpus 8其中X.py换成任意一个配置文件名。等价地你也可以在仓库根目录下执行python tools/lazyconfig_train_net.py \ --config-file projects/Rethinking-BatchNorm/configs/mask_rcnn_BNhead.py \ --num-gpus 8注意两点必须使用tools/lazyconfig_train_net.py而非tools/plain_train_net.py这些配置是 Python 形式的 LazyConfig通过LazyConfig与L()惰性实例化与 YACS 风格的 yaml 配置不兼容--num-gpus 8并非可选SyncBN、batch shuffle 等实验依赖torch.distributed多卡通信单卡运行会导致统计口径与论文不一致。--num-gpus会被 engine/launch.py 解析并启动 DDP 训练。所有配置都通过get_config继承仓库根目录 configs/common 下的公共组件详见下文公共配置继承小节因此每个文件都非常精简——这正是 LazyConfig 设计的复用模式。三、Mask R-CNN 系列head 中的 BatchNorm 实验3.1 mask_rcnn_BNhead.py — head 中加入 BatchNorm对应论文 Table 3该配置是后续三个 Mask R-CNN 变体的基类配置全文如下from detectron2.model_zoo import get_config model get_config(common/models/mask_rcnn_fpn.py).model model.backbone.bottom_up.freeze_at 2 model.roi_heads.box_head.conv_norm model.roi_heads.mask_head.conv_norm BN # 4conv1fc head model.roi_heads.box_head.conv_dims [256, 256, 256, 256] model.roi_heads.box_head.fc_dims [1024] dataloader get_config(common/data/coco.py).dataloader lr_multiplier get_config(common/coco_schedule.py).lr_multiplier_3x optimizer get_config(common/optim.py).SGD train get_config(common/train.py).train train.init_checkpoint detectron2://ImageNetPretrained/MSRA/R-50.pkl train.max_iter 270000 # 3x for batchsize 16逐项解读配置项取值含义model.backbone.bottom_up.freeze_at 22冻结 ResNet 前两个 stagestem res2的参数只更新后面层。这与 common/models/mask_rcnn_fpn.py 中的默认设置一致box_head.conv_norm BN字符串BN在 box head 的 4 个卷积层后插入 BatchNorm。该字符串会被 box_head.py 中的get_norm(conv_norm, conv_dim)解析为nn.BatchNorm2dmask_head.conv_norm BN字符串BN同理在 mask head 的卷积层后插入 BN见 mask_head.py 中conv_norm参数box_head.conv_dims [256,256,256,256]、fc_dims [1024]—4 层 256 通道卷积 1 层 1024 维全连接即注释中的 4conv1fc 检测头结构train.init_checkpointR-50.pkl从 Detectron2 的 ImageNet 预训练权重MSRA R-50初始化train.max_iter 270000270k3x 训练计划约 37 epochCOCO对应 common/coco_schedule.py 中的lr_multiplier_3x注释明确说明该数值基于 total batch size 16源码佐证在FastRCNNConvFCHead的构造中每个卷积层通过get_norm(conv_norm, conv_dim)生成归一化层且当指定了 norm 时卷积层不设 biasbiasnot conv_norm因为 BN 的 affine 变换会吸收 bias——见 box_head.py。这就是把 head 从无归一化改成带 BN时只有一行配置的原因检测头的卷积层早已支持 norm 参数。3.2 mask_rcnn_BNhead_batch_stats.py — 推理期使用 batch 统计量from torch.nn import BatchNorm2d from torch.nn import functional as F class BatchNormBatchStat(BatchNorm2d): BN that uses batch stat in inference def forward(self, input): if self.training: return super().forward(input) return F.batch_norm(input, None, None, self.weight, self.bias, True, 1.0, self.eps) # After training with the base config, its sufficient to load its model with # this config only for inference -- because the training-time behavior is identical. from .mask_rcnn_BNhead import model, dataloader, lr_multiplier, optimizer, train model.roi_heads.box_head.conv_norm model.roi_heads.mask_head.conv_norm BatchNormBatchStat关键设计它定义了一个继承BatchNorm2d的BatchNormBatchStat只在推理期改变行为训练时走标准 BNsuper().forward推理时调用F.batch_norm(input, None, None, ...)传入的running_mean/running_var均为None且trainingTrueF.batch_norm的第四个布尔参数即强制用当前 batch 的统计量而非滑动平均因此训练阶段与基类配置完全一致。只需用基类配置mask_rcnn_BNhead.py训练好的模型再套上本配置做推理即可——无需重新训练这是文件末尾注释强调的要点在 COCO 评测场景下若推理 batch 足够大用 batch 统计量替代全局统计量能反映batch 大小对 BN 的影响对应论文 Table 3 的消融设计。3.3 mask_rcnn_BNhead_shuffle.py — 跨 GPU 打乱 head 输入对应论文 Figure 9 / Table 6这个配置在源码层面最复杂它用运行时动态构造子类的方式给 head 的输入做跨卡随机打乱再还原import math import torch import torch.distributed as dist from detectron2.modeling.roi_heads import FastRCNNConvFCHead, MaskRCNNConvUpsampleHead from detectron2.utils import comm from fvcore.nn.distributed import differentiable_all_gather def concat_all_gather(input): bs_int input.shape[0] size_list comm.all_gather(bs_int) max_size max(size_list) max_shape (max_size,) input.shape[1:] padded_input input.new_zeros(max_shape) padded_input[:bs_int] input all_inputs differentiable_all_gather(padded_input) inputs [x[:sz] for sz, x in zip(size_list, all_inputs)] return inputs, size_list def batch_shuffle(x): # gather from all gpus batch_size_this x.shape[0] all_xs, batch_size_all concat_all_gather(x) all_xs_concat torch.cat(all_xs, dim0) total_bs sum(batch_size_all) rank dist.get_rank() assert batch_size_all[rank] batch_size_this idx_range (sum(batch_size_all[:rank]), sum(batch_size_all[: rank 1])) # random shuffle index idx_shuffle torch.randperm(total_bs, devicex.device) # broadcast to all gpus dist.broadcast(idx_shuffle, src0) # index for restoring idx_unshuffle torch.argsort(idx_shuffle) # shuffled index for this gpu splits torch.split(idx_shuffle, math.ceil(total_bs / dist.get_world_size())) if len(splits) rank: idx_this splits[rank] else: idx_this idx_shuffle.new_zeros([0]) return all_xs_concat[idx_this], idx_unshuffle[idx_range[0] : idx_range[1]] def batch_unshuffle(x, idx_unshuffle): all_x, _ concat_all_gather(x) x_gather torch.cat(all_x, dim0) return x_gather[idx_unshuffle] def wrap_shuffle(module_type, method): def new_method(self, x): if self.training: x, idx batch_shuffle(x) x getattr(module_type, method)(self, x) if self.training: x batch_unshuffle(x, idx) return x return type(module_type.__name__ WithShuffle, (module_type,), {method: new_method}) from .mask_rcnn_BNhead import model, dataloader, lr_multiplier, optimizer, train model.roi_heads.box_head._target_ wrap_shuffle(FastRCNNConvFCHead, forward) model.roi_heads.mask_head._target_ wrap_shuffle(MaskRCNNConvUpsampleHead, layers)实现原理分三步concat_all_gather用fvcore.nn.distributed.differentiable_all_gather做可微的全量收集区别于torch.distributed.all_gather的不可微版本并对每个 GPU 上的输入先 padding 到统一形状再收集、收集后按各自真实 batch 截断——保证反向传播能正确回传梯度batch_shuffle把所有 GPU 的特征拼成一个大 batch用torch.randperm生成全局随机索引通过dist.broadcast(..., src0)保证所有 rank 拿到同一份乱序索引随后每个 GPU 领取自己那份打乱后的数据同时返回idx_unshuffle逆置换索引用于后续还原batch_unshufflehead 计算完之后再次跨卡收集结果用之前保存的逆索引把样本放回原始位置保证 loss 计算时每个样本的预测与标签对齐。wrap_shuffle是点睛之笔它不是手写一个新的 head 类而是用type()在运行时基于FastRCNNConvFCHead/MaskRCNNConvUpsampleHead动态生成一个XXXWithShuffle子类仅覆盖forwardbox head或layersmask head方法在调用原始方法前后插入 shuffle / unshuffle。这样只需要改 LazyConfig 的_target_就能让已实例化的 head 结构整体替换。实验意义打乱 head 输入等价于让每个 BN 层看到的 batch不再局限于本卡的样本而近似于更大的、随机采样的跨卡 batch用于隔离batch 内样本相关性对 BN 统计量的影响——对应论文 Figure 9 与 Table 6 的讨论。3.4 mask_rcnn_SyncBNhead.py — head 中使用 SyncBN对应论文 Table 6from .mask_rcnn_BNhead import model, dataloader, lr_multiplier, optimizer, train model.roi_heads.box_head.conv_norm model.roi_heads.mask_head.conv_norm SyncBN在基类配置之上仅改一行把conv_norm从BN换成SyncBN。get_norm会将其解析为NaiveSyncBatchNorm见 layers/batch_norm.py其内部通过dist.all_reduce汇总各卡统计量后统一归一化。README 指出该配置可匹配论文 Table 6 的结果。对比小结三份配置的差异都在 head 的归一化策略上backbone 与训练计划完全一致配置head 归一化方式训练期统计范围mask_rcnn_BNhead.py单卡 BN本卡 mini-batchmask_rcnn_BNhead_batch_stats.py单卡 BN推理期用 batch 统计本卡 mini-batchmask_rcnn_BNhead_shuffle.py单卡 BN 跨卡输入打乱打乱后的跨卡混合 batchmask_rcnn_SyncBNhead.pySyncBN跨卡全局 batch四、RetinaNet 系列特征金字塔多域下的 BN 实验RetinaNet 与 Mask R-CNN 的关键差异在于5 个金字塔特征层p3–p7共享同一个检测头。同一组 BN 层被 5 个输入域轮流调用于是出现BN 的 batch 统计到底按哪个域算的问题——这正是论文 Table 5 的实验主题。4.1 retinanet_SyncBNhead.py — head 中使用 SyncBN对应论文 Table 5 row 3from detectron2.model_zoo import get_config from torch import nn model get_config(common/models/retinanet.py).model model.backbone.bottom_up.freeze_at 2 # The head will overwrite string SyncBN to use domain-specific BN, so we # provide a class here to use shared BN in training. model.head.norm nn.SyncBatchNorm2d dataloader get_config(common/data/coco.py).dataloader lr_multiplier get_config(common/coco_schedule.py).lr_multiplier_3x optimizer get_config(common/optim.py).SGD train get_config(common/train.py).train optimizer.lr 0.01 train.init_checkpoint detectron2://ImageNetPretrained/MSRA/R-50.pkl train.max_iter 270000 # 3x for batchsize 16这里有一个容易踩坑的细节注释明确指出——如果给model.head.norm传字符串SyncBNRetinaNetHead会劫持它并自动改成域特定domain-specificBN。看源码 retinanet.pyif norm BN or norm SyncBN: logger.info( fUsing domain-specific {norm} in RetinaNetHead with len{self._num_features}. ) bn_class nn.BatchNorm2d if norm BN else nn.SyncBatchNorm def norm(c): return CycleBatchNormList( lengthself._num_features, bn_classbn_class, num_featuresc )也就是说RetinaNetHead认为在共享 head 场景下使用普通sharedBN 效果不佳源码中遇到 shared BN 还会给出 warningShared BatchNorm may not work well in RetinaNetHead因此字符串BN/SyncBN一律被替换为CycleBatchNormList域特定 BN见第六节。要强行使用 shared BN就必须像本配置这样直接传类对象nn.SyncBatchNorm2d非字符串不会被字符串分支捕获会走get_norm正常解析。其余要点与 Mask R-CNN 系列一样冻结 backbone 前两层freeze_at 2、使用 COCO 数据与 3x 计划、270k 迭代optimizer.lr 0.01单独指定了 SGD 初始学习率RetinaNet 训练常用 0.01 配合 batch size 16该配置是straightforward的 SyncBN-in-head 实现README 说明其匹配论文 Table 5 的 row 3。4.2 retinanet_SyncBNhead_SharedTraining.py — 5 个特征层共享归一化对应论文 Table 5 row 1from typing import List import torch from torch import Tensor, nn from detectron2.modeling.meta_arch.retinanet import RetinaNetHead def apply_sequential(inputs, modules): for mod in modules: if isinstance(mod, (nn.BatchNorm2d, nn.SyncBatchNorm)): # for BN layer, normalize all inputs together shapes [i.shape for i in inputs] spatial_sizes [s[2] * s[3] for s in shapes] x [i.flatten(2) for i in inputs] x torch.cat(x, dim2).unsqueeze(3) x mod(x).split(spatial_sizes, dim2) inputs [i.view(s) for s, i in zip(shapes, x)] else: inputs [mod(i) for i in inputs] return inputs class RetinaNetHead_SharedTrainingBN(RetinaNetHead): def forward(self, features: List[Tensor]): logits apply_sequential(features, list(self.cls_subnet) [self.cls_score]) bbox_reg apply_sequential(features, list(self.bbox_subnet) [self.bbox_pred]) return logits, bbox_reg from .retinanet_SyncBNhead import model, dataloader, lr_multiplier, optimizer, train model.head._target_ RetinaNetHead_SharedTrainingBN核心是apply_sequential函数它逐个遍历 head 的子模块遇到 BN 层时把 5 个特征层先各自flatten(2)展平、沿通道维torch.cat拼接成一个大矩阵再统一送入同一个 BN 层归一化最后按各层的空间尺寸split还原。这样一个 BN 层的 batch 统计量同时聚合了 5 个金字塔层、所有空间位置的样本——即共享归一化统计的训练方式。非 BN 层卷积、激活等则按常规逐层逐个特征图计算。RetinaNetHead_SharedTrainingBN继承RetinaNetHead并重写forward把 cls 子网和 bbox 子网含各自的最终输出层都换成apply_sequential处理。配置中通过model.head._target_ RetinaNetHead_SharedTrainingBN替换 LazyConfig 的惰性目标类即可继承自上一节的retinanet_SyncBNhead.py的全部训练设置。README 说明该配置匹配论文 Table 5 的 row 1。两种 RetinaNet 变体的差异一句话总结retinanet_SyncBNhead.pyhead 中每层 BN 的统计来自单个特征层域特定 BN跨卡 Sync 域内统计传nn.SyncBatchNorm2d类以绕过字符串劫持retinanet_SyncBNhead_SharedTraining.pyhead 中每层 BN 的统计来自全部 5 个特征层拼接后的大 batch共享 BN。五、域特定统计量评测retinanet-eval-domain-specific.py对应论文 Table 5 row 4 / row 2训练完成后论文还讨论了一个评测问题RetinaNet 的 head 被 5 个特征域复用推理期每个域应该用自己域的滑动统计量。脚本 retinanet-eval-domain-specific.py 在加载 checkpoint 后重新计算域特定统计量再评测./retinanet-eval-domain-specific.py checkpoint.pth脚本核心逻辑from fvcore.nn.precise_bn import update_bn_stats from detectron2.checkpoint import DetectionCheckpointer from detectron2.config import LazyConfig, instantiate from detectron2.evaluation import inference_on_dataset from detectron2.layers import CycleBatchNormList ... cfg LazyConfig.load_rel(configs/retinanet_SyncBNhead.py) model cfg.model model.head.norm lambda c: CycleBatchNormList(len(model.head_in_features), num_featuresc) model instantiate(model) model.cuda() DetectionCheckpointer(model).load(checkpoint) cfg.dataloader.train.total_batch_size 8 with EventStorage(), torch.no_grad(): update_bn_stats(model, instantiate(cfg.dataloader.train), 500) inference_on_dataset(model, ...)要点强制域特定 BN即使训练用的是 shared BN 配置上一节的两个变体评测时统一把model.head.norm覆盖为CycleBatchNormList(length5, ...)——为 5 个金字塔层各维护一份 BN 统计len(model.head_in_features)即 p3–p7 的数量 5重算统计量fvcore.nn.precise_bn.update_bn_stats(model, dataloader, 500)在 500 个训练 batch 上以前向无梯度模式重新累计每层每个域的均值/方差替代训练时的滑动平均使统计量与推理 batch 口径一致小 batch 重算total_batch_size 8用于控制重算统计量时的 batch 规模。README 说明对上述两个 RetinaNet 配置训练出的模型运行该脚本结果可分别匹配论文 Table 5 的 row 4 与 row 2——这正是训练用共享/域特定 BN评测统一用域特定统计量的组合带来的精度提升。六、底层机制CycleBatchNormList 与 get_norm6.1 CycleBatchNormList —— 域特定 BN 的实现上文反复出现的CycleBatchNormList定义在 detectron2/layers/batch_norm.py其文档字符串明确写着 Implement domain-specific BatchNorm by cycling并直接引用论文 Sec 5.2class CycleBatchNormList(nn.ModuleList): Implement domain-specific BatchNorm by cycling. When a BatchNorm layer is used for multiple input domains or input features, it might need to maintain a separate test-time statistics for each domain. See Sec 5.2 in :paper:rethinking-batchnorm. This module implements it by using N separate BN layers and it cycles through them every time a forward() is called. NOTE: The caller of this module MUST guarantee to always call this module by multiple of N times. Otherwise its test-time statistics will be incorrect. def __init__(self, length: int, bn_classnn.BatchNorm2d, **kwargs): self._affine kwargs.pop(affine, True) super().__init__([bn_class(**kwargs, affineFalse) for k in range(length)]) if self._affine: # shared affine, domain-specific BN channels self[0].num_features self.weight nn.Parameter(torch.ones(channels)) self.bias nn.Parameter(torch.zeros(channels)) self._pos 0 def forward(self, x): ret selfself._pos self._pos (self._pos 1) % len(self) if self._affine: w self.weight.reshape(1, -1, 1, 1) b self.bias.reshape(1, -1, 1, 1) ...设计要点N 个独立 BN 子层 循环调度第 k 次调用使用第k % N个子层从而为第 k 个输入域维护独立的 running_mean / running_var共享 affine 参数weight/bias不放在子 BN 里子层affineFalse而是由外层共享一个可学习参数——即 shared affine, domain-specific BN所有域共享缩放平移但各自维护统计量调用纪律注释特别警告调用方必须保证调用次数是 N 的整数倍否则测试期统计量会错位——在 RetinaNet 中RetinaNetHead每次 forward 恰好处理_num_features5个特征层天然满足该约束。6.2 head 中 norm 的解析链路Mask R-CNNbox / mask head 的conv_norm参数传入get_norm(conv_norm, conv_dim)生成归一化层见 box_head.py、mask_head.py。get_norm定义于 layers/batch_norm.py支持BN、SyncBN、FrozenBN、GN等字符串以及任意 callableRetinaNetRetinaNetHead在构造函数中特殊处理norm BN or SyncBN拦截字符串并替换为CycleBatchNormList见 retinanet.py这就是域特定 BN自动化的位置传类对象则绕过该分支。6.3 公共配置继承本项目的 6 份配置几乎都通过get_config(common/...)复用仓库根目录 configs/common 下的公共组件公共模块提供的对象说明common/models/mask_rcnn_fpn.pymodelMask R-CNN FPN R-50 模型骨架common/models/retinanet.pymodelRetinaNet 模型骨架head_in_features[p3,p4,p5,p6,p7]common/data/coco.pydataloaderCOCO 数据加载器common/coco_schedule.pylr_multiplier_3x3x 学习率调度270k 迭代common/optim.pySGDSGD 优化器momentum、weight decay 默认值common/train.pytrain训练超参output_dir、checkpoint 周期、eval 周期、AMP/DDP 选项等这种配置组合composition模式是 LazyConfig 相对 yaml 的核心优势每个实验只需写与基线不同的差异行并可用from .base import ...或get_config任意复用。七、复现实验时的注意事项多卡是硬前提mask_rcnn_BNhead_shuffle.py依赖dist.broadcast/differentiable_all_gathermask_rcnn_SyncBNhead.py与retinanet_SyncBNhead.py依赖 SyncBN 的跨卡 all-reduce。单卡或未初始化 DDP 环境运行会报错或产生与论文不一致的统计口径batch size 对齐train.max_iter 270000的注释强调 3x for batchsize 16调整 total batch size 时需要同步调整迭代数与学习率RetinaNet 配置中optimizer.lr 0.01同样以 batch size 16 为前提评测与训练分离mask_rcnn_BNhead_batch_stats.py无需重新训练训练行为与基类一致retinanet-eval-domain-specific.py接受任意 checkpoint 路径重算域特定统计量后再评测两处都是训练配置 评测配置解耦的典型用法域特定 BN 的自动化陷阱给RetinaNetHead传字符串BN/SyncBN会被自动替换为CycleBatchNormList需要真正的 shared BN 时务必传类对象如nn.SyncBatchNorm2d否则实验结果与预期不一致。结语从这份精简的 README 出发可以看到 batch 一词在 BatchNorm 语境下具有多层含义单卡 batchmask_rcnn_BNhead、跨卡 batchSyncBN、跨卡混合 batchshuffle、跨特征域 batchshared training BN以及域特定统计量CycleBatchNormList。本文涉及的 6 份 LazyConfig 与 1 个评测脚本完整覆盖了这些语义其对应的论文实验结果分别为 Mask R-CNN 的 Table 3 / Figure 9 / Table 6 与 RetinaNet 的 Table 5row 1–4。如需在自有数据上复现可对照 配置文件 逐行修改公共配置数据、迭代数、学习率并保持多卡训练与评测统计口径的一致性。【免费下载链接】detectron2Detectron2 is a platform for object detection, segmentation and other visual recognition tasks.项目地址: https://gitcode.com/GitHub_Trending/de/detectron2创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考