2026-06-14-xnet-groupnorm-design.md 2.6 KB

XNet GroupNorm 替换设计

1. 目标

将当前 XNet2d 的 2D 卷积主路径从 BatchNorm2d 切换到更适合 2D 图像分割、且对 batch size 更稳健的 GroupNorm,目标是:

  1. 保持当前 2D 分割主线可训练。
  2. 避免深层特征在 batch_size=1 或极小空间尺寸下因 BatchNorm2d 统计约束报错。
  3. 采用工程上常见、对 2D 医学分割友好的默认方案。

2. 推荐方案

当前主线采用:

CNN / 卷积块 -> GroupNorm

不修改 lib/modules/lib_mamba/vmamba.py 内部的 LayerNorm / BatchNorm 逻辑,本轮只覆盖 XNet2d 主路径直接使用的 2D 卷积归一化层。

3. 替换范围

3.1 代码

主要修改:

  • lib/modules/layers_2d.py

受影响调用者:

  • lib/modules/xnet_2d.py
  • 可能还有同目录下复用 Conv2dBN / DWConv2dBNReLU 的模块

3.2 具体做法

  1. 不保留 Conv2dBN 兼容名,直接引入 Conv2dGN
  2. Conv2dBN 的调用点统一替换为 Conv2dGN
  3. DWConv2dBNReLU 重命名为 DWConv2dGNReLU,并将内部两个 BatchNorm2d 改为 GroupNorm
  4. 保留现有 bn_weight_init 参数语义,但内部作用到 GroupNorm.weight;如有必要,后续可再统一改名为更中性的参数名。

4. GroupNorm 策略

默认采用自动组数选择:

  1. 优先尝试 8 组。
  2. 若通道数不能被 8 整除,则递减选择能整除的组数。
  3. 最差回退到 1 组。

这样可以兼顾:

  • 常见通道数如 8, 16, 24, 32, 64, 128, 192
  • 小通道层与深层大通道层
  • 不额外强制用户手动配组数

5. 保守边界

本轮不做以下改动:

  1. 不统一替换 lib/modules/lib_mamba/vmamba.py 中的 norm。
  2. 不引入新的 YAML 配置项让用户在 BN/GN 间切换。
  3. 不扩展为多种 norm 可配置框架,本轮直接收敛到 GN 命名与实现。

原因是本次目标是优先得到“2D 分割主线稳定、主流、可训练”的默认行为,而不是构建归一化实验框架。

6. 验证目标

需要验证三类证据:

  1. pytest -q tests/test_xnet_2d.py 通过。
  2. 最小训练烟测通过,且可在 batch_size=164x64 条件下至少跑通 1 个 epoch。
  3. 代码搜索确认 layers_2d.py 中主线 2D 卷积归一化已不再使用 BatchNorm2d

7. 风险与预期

风险:

  1. 训练数值行为会和旧 BN 版本不同。
  2. 已有历史实验结果与新结果不可直接横向对比。

预期收益:

  1. 小 batch 与不固定 batch 更稳。
  2. 2D 医学分割主线更符合常见工程实践。
  3. 烟测不再受深层 1x1 + BN 统计限制影响。