feat(layers): 将2D卷积归一化从BatchNorm切换到GroupNorm

- 将Conv2dBN重命名为Conv2dGN，使用GroupNorm替代BatchNorm
- 将DWConv2dBNReLU重命名为DWConv2dGNReLU，内部BatchNorm改为GroupNorm
- 在PatchMerging2d等模块中同步更新卷积归一化层引用
- 实现自动组数选择策略，优先8组，不能整除则递减到可整除，最差回退到1组

docs: 更新文档描述2D卷积使用GroupNorm

- 修改README.md中关于2D卷积主路径使用GroupNorm的说明
- 更新XNet_method.md文档中的架构描述
- 添加XNet GroupNorm替换设计文档和迁移计划文档

test: 增加GroupNorm相关测试确保无BatchNorm残留

- 添加测试验证XNet2d主路径不再使用BatchNorm2d
- 增加batch_size=1训练测试避免BatchNorm统计错误
- 更新模块导出列表同步新命名约定

README.md

@@ -12,8 +12,9 @@ CNN-Wavelet-VMamba encoder + plain U-Net skip decoder + segmentation head
 
 1. encoder 使用 `XTEB2d`，融合 local、wavelet 和 VMamba-style `SS2D` 三个分支。
 2. decoder 已不再使用历史版本中的 diagonal guide / X-shaped guide path，而是普通 U-Net 同尺度 skip 融合。
-3. decoder 可选启用 `frequency refine`。
-4. 训练主链通过 `tools/run_optimized_supervised.sh` 或 `python tools/train.py --config ...` 启动。
+3. 2D 卷积主路径默认使用 `GroupNorm`，以适配大小 batch 都可能出现的分割训练场景。
+4. decoder 可选启用 `frequency refine`。
+5. 训练主链通过 `tools/run_optimized_supervised.sh` 或 `python tools/train.py --config ...` 启动。
 
 建议先阅读以下文档：
 

docs/method/XNet_method.md

@@ -16,8 +16,9 @@ CNN-Wavelet-VMamba encoder + plain U-Net skip decoder + segmentation head
 1. encoder 使用 `XTEB2d`，在每个尺度内融合局部纹理、小波频率信息和 VMamba-style SS2D 全局建模。
 2. bottleneck 继续使用 `XTEB2d` 加强最深层语义表征。
 3. decoder 使用普通 U-Net 式逐级上采样和同尺度 skip 融合。
-4. decoder block 内可选使用频率细化模块，帮助恢复低对比度边界和局部细节。
-5. segmentation head 将 `H/4 x W/4` 的 decoder 输出上采样回输入分辨率。
+4. 2D 卷积主路径默认采用 `GroupNorm`，减少对 batch size 的依赖。
+5. decoder block 内可选使用频率细化模块，帮助恢复低对比度边界和局部细节。
+6. segmentation head 将 `H/4 x W/4` 的 decoder 输出上采样回输入分辨率。
 
 旧版设计中的 `XGuideProjector2d`、`XGuideModulation2d`、diagonal guide path、`guide_mode` 配置字段以及 `guides` 输出字段都已从当前主线代码中移除。
 

docs/superpowers/plans/2026-06-14-xnet-conv2dgn-migration.md

@@ -0,0 +1,57 @@
+# XNet Conv2dGN Migration Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** 将 `XNet2d` 主路径的 2D 卷积归一化从 BatchNorm 明确迁移为 GroupNorm，并统一命名为 `Conv2dGN` / `DWConv2dGNReLU`。
+
+**Architecture:** 只替换 `lib/modules/layers_2d.py` 及其主线调用者中的 2D 卷积归一化层，不触碰 `lib/modules/lib_mamba/vmamba.py` 内部的 LayerNorm/BatchNorm 逻辑。采用自动组数选择的 `GroupNorm`，优先保证 2D 分割训练在大小 batch 下都稳定。
+
+**Tech Stack:** PyTorch、GroupNorm、pytest、现有 XNet2d 训练与烟测入口
+
+---
+
+### Task 1: 先收紧测试约束
+
+**Files:**
+- Modify: `tests/test_xnet_2d.py`
+
+- [ ] 增加一条针对归一化实现的测试，验证 `XNet2d` 主路径中不再实例化 `nn.BatchNorm2d`。
+- [ ] 先运行 `pytest -q tests/test_xnet_2d.py`，确认新测试先失败，再做实现替换。
+
+### Task 2: 实现 `Conv2dGN` 与 `DWConv2dGNReLU`
+
+**Files:**
+- Modify: `lib/modules/layers_2d.py`
+
+- [ ] 增加自动 group 数选择 helper，优先 `8` 组，不能整除则递减到可整除，最差回退 `1`。
+- [ ] 新增 `Conv2dGN`，用 `Conv2d + GroupNorm` 替代 `Conv2dBN`。
+- [ ] 将 `DWConv2dBNReLU` 重命名并实现为 `DWConv2dGNReLU`。
+- [ ] `PatchMerging2d` 等直接依赖层同步改用 `Conv2dGN`。
+
+### Task 3: 替换 XNet 调用点与模块导出
+
+**Files:**
+- Modify: `lib/modules/xnet_2d.py`
+- Modify: `lib/modules/__init__.py`
+
+- [ ] 将 `xnet_2d.py` 中所有 `Conv2dBN` 引用替换为 `Conv2dGN`。
+- [ ] 将 `__init__.py` 中导出名称同步改成 `Conv2dGN` / `DWConv2dGNReLU`，不再保留旧 BN 名称。
+
+### Task 4: 运行测试并做 batch_size=1 烟测
+
+**Files:**
+- Modify as needed: `tests/test_xnet_2d.py`
+
+- [ ] 运行 `pytest -q tests/test_xnet_2d.py`，确认通过。
+- [ ] 运行 1 epoch 最小烟测，目标配置为 `batch_size=1`、`64x64`、`validation.enabled=false`。
+- [ ] 若烟测仍失败，先确认是否还有残留 `BatchNorm2d`，而不是直接扩大 batch 规避。
+
+### Task 5: 更新当前主线文档
+
+**Files:**
+- Modify: `README.md`
+- Modify: `docs/method/XNet_method.md`
+- Modify: `docs/training/当前项目详解与纯文本架构流程图.md`
+
+- [ ] 将当前主线路径中的 norm 描述改为 GroupNorm。
+- [ ] 在方法和训练文档中简要说明：当前 2D 卷积主路径使用 GroupNorm，以支持更稳健的小 batch 与不固定 batch 训练。

docs/superpowers/specs/2026-06-14-xnet-groupnorm-design.md

@@ -0,0 +1,84 @@
+# XNet GroupNorm 替换设计
+
+## 1. 目标
+
+将当前 `XNet2d` 的 2D 卷积主路径从 `BatchNorm2d` 切换到更适合 2D 图像分割、且对 batch size 更稳健的 `GroupNorm`，目标是：
+
+1. 保持当前 2D 分割主线可训练。
+2. 避免深层特征在 `batch_size=1` 或极小空间尺寸下因 `BatchNorm2d` 统计约束报错。
+3. 采用工程上常见、对 2D 医学分割友好的默认方案。
+
+## 2. 推荐方案
+
+当前主线采用：
+
+```text
+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=1`、`64x64` 条件下至少跑通 1 个 epoch。
+3. 代码搜索确认 `layers_2d.py` 中主线 2D 卷积归一化已不再使用 `BatchNorm2d`。
+
+## 7. 风险与预期
+
+风险：
+
+1. 训练数值行为会和旧 BN 版本不同。
+2. 已有历史实验结果与新结果不可直接横向对比。
+
+预期收益：
+
+1. 小 batch 与不固定 batch 更稳。
+2. 2D 医学分割主线更符合常见工程实践。
+3. 烟测不再受深层 `1x1` + BN 统计限制影响。

docs/training/当前项目详解与纯文本架构流程图.md

@@ -47,6 +47,7 @@ outputs["seg_logits"]
 ```text
 用 XNet2d 在 BUSI / DDTI / TN3K / TG3K 等 2D 超声数据集上做全监督分割训练。
 XNet2d 的 encoder 用 local + wavelet + VMamba-style SS2D 三分支建模，
+卷积主路径默认使用 GroupNorm，
 decoder 用普通 U-Net 同尺度 skip + 频率细化恢复 mask。
 ```
 
@@ -732,7 +733,7 @@ output: Y [B,C,H,W]
 ```text
 X
 │
-├─ pre_norm: 1x1 Conv2dBN
+├─ pre_norm: 1x1 Conv2dGN
 │
 ├─ Local branch
 │   ├─ DWConv3x3 + PWConv1x1

lib/modules/__init__.py

@@ -1,7 +1,7 @@
 from .layers_2d import (
     BNLinear1d,
-    Conv2dBN,
-    DWConv2dBNReLU,
+    Conv2dGN,
+    DWConv2dGNReLU,
     DropPath,
     FFN2d,
     PatchMerging2d,
@@ -23,8 +23,8 @@ from .xnet_2d import (
 
 __all__ = [
     "BNLinear1d",
-    "Conv2dBN",
-    "DWConv2dBNReLU",
+    "Conv2dGN",
+    "DWConv2dGNReLU",
     "DropPath",
     "FFN2d",
     "PatchMerging2d",

lib/modules/layers_2d.py

@@ -55,22 +55,29 @@ class FFN2d(nn.Module):
         return self.mlp(x)
 
 
-class BNLinear1d(nn.Sequential):
+class BNLinear1d(nn.Sequential):
     def __init__(self, in_features, out_features, bias=True, std=0.02):
         bn = nn.BatchNorm1d(in_features)
         linear = nn.Linear(in_features, out_features, bias=bias)
         trunc_normal_(linear.weight, std=std)
-        if bias:
-            nn.init.constant_(linear.bias, 0)
-        super().__init__(OrderedDict([("bn", bn), ("linear", linear)]))
-
-
-class Conv2dBN(nn.Sequential):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size=1,
+        if bias:
+            nn.init.constant_(linear.bias, 0)
+        super().__init__(OrderedDict([("bn", bn), ("linear", linear)]))
+
+
+def _resolve_group_count(num_channels: int, preferred_groups: int = 8) -> int:
+    groups = min(preferred_groups, num_channels)
+    while groups > 1 and num_channels % groups != 0:
+        groups -= 1
+    return groups
+
+
+class Conv2dGN(nn.Sequential):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=1,
         stride=1,
         padding=0,
         dilation=1,
@@ -86,18 +93,18 @@ class Conv2dBN(nn.Sequential):
             dilation,
             groups,
             bias=False,
-        )
-        bn = nn.BatchNorm2d(out_channels)
-        nn.init.constant_(bn.weight, bn_weight_init)
-        nn.init.constant_(bn.bias, 0)
-        super().__init__(OrderedDict([("conv", conv), ("bn", bn)]))
-
-
-class DWConv2dBNReLU(nn.Sequential):
-    def __init__(self, in_channels, out_channels, kernel_size=3, bn_weight_init=1.0):
-        super().__init__(
-            OrderedDict(
-                [
+        )
+        gn = nn.GroupNorm(_resolve_group_count(out_channels), out_channels)
+        nn.init.constant_(gn.weight, bn_weight_init)
+        nn.init.constant_(gn.bias, 0)
+        super().__init__(OrderedDict([("conv", conv), ("gn", gn)]))
+
+
+class DWConv2dGNReLU(nn.Sequential):
+    def __init__(self, in_channels, out_channels, kernel_size=3, bn_weight_init=1.0):
+        super().__init__(
+            OrderedDict(
+                [
                     (
                         "dwconv3x3",
                         nn.Conv2d(
@@ -106,43 +113,43 @@ class DWConv2dBNReLU(nn.Sequential):
                             kernel_size,
                             1,
                             kernel_size // 2,
-                            groups=in_channels,
-                            bias=False,
-                        ),
-                    ),
-                    ("bn1", nn.BatchNorm2d(in_channels)),
-                    ("relu", nn.ReLU(inplace=True)),
-                    (
-                        "dwconv1x1",
-                        nn.Conv2d(
+                            groups=in_channels,
+                            bias=False,
+                        ),
+                    ),
+                    ("gn1", nn.GroupNorm(_resolve_group_count(in_channels), in_channels)),
+                    ("relu", nn.ReLU(inplace=True)),
+                    (
+                        "dwconv1x1",
+                        nn.Conv2d(
                             in_channels,
                             out_channels,
                             1,
                             1,
                             0,
-                            groups=in_channels,
-                            bias=False,
-                        ),
-                    ),
-                    ("bn2", nn.BatchNorm2d(out_channels)),
-                ]
-            )
-        )
-        for bn_name in ["bn1", "bn2"]:
-            bn = getattr(self, bn_name)
-            nn.init.constant_(bn.weight, bn_weight_init)
-            nn.init.constant_(bn.bias, 0)
-
-
-class PatchMerging2d(nn.Module):
-    def __init__(self, dim, out_dim):
-        super().__init__()
-        hidden_dim = int(dim * 4)
-        self.conv1 = Conv2dBN(dim, hidden_dim, 1, 1, 0)
-        self.act = nn.ReLU(inplace=True)
-        self.conv2 = Conv2dBN(hidden_dim, hidden_dim, 3, 2, 1, groups=hidden_dim)
-        self.se = SqueezeExcite(hidden_dim, rd_ratio=0.25)
-        self.conv3 = Conv2dBN(hidden_dim, out_dim, 1, 1, 0)
+                            groups=in_channels,
+                            bias=False,
+                        ),
+                    ),
+                    ("gn2", nn.GroupNorm(_resolve_group_count(out_channels), out_channels)),
+                ]
+            )
+        )
+        for gn_name in ["gn1", "gn2"]:
+            gn = getattr(self, gn_name)
+            nn.init.constant_(gn.weight, bn_weight_init)
+            nn.init.constant_(gn.bias, 0)
+
+
+class PatchMerging2d(nn.Module):
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        hidden_dim = int(dim * 4)
+        self.conv1 = Conv2dGN(dim, hidden_dim, 1, 1, 0)
+        self.act = nn.ReLU(inplace=True)
+        self.conv2 = Conv2dGN(hidden_dim, hidden_dim, 3, 2, 1, groups=hidden_dim)
+        self.se = SqueezeExcite(hidden_dim, rd_ratio=0.25)
+        self.conv3 = Conv2dGN(hidden_dim, out_dim, 1, 1, 0)
 
     def forward(self, x):
         return self.conv3(self.se(self.act(self.conv2(self.act(self.conv1(x))))))

lib/modules/xnet_2d.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .layers_2d import Conv2dBN
+from .layers_2d import Conv2dGN
 from .lib_mamba.vmamba import SS2D as VMambaSS2D
 
 
@@ -16,13 +16,13 @@ class XNetStem2d(nn.Module):
     def __init__(self, in_channels: int, stem_channels: int, out_channels: int) -> None:
         super().__init__()
         self.block = nn.Sequential(
-            Conv2dBN(in_channels, stem_channels, 3, 2, 1),
+            Conv2dGN(in_channels, stem_channels, 3, 2, 1),
             nn.ReLU(inplace=True),
-            Conv2dBN(stem_channels, stem_channels, 3, 1, 1, groups=stem_channels),
+            Conv2dGN(stem_channels, stem_channels, 3, 1, 1, groups=stem_channels),
             nn.ReLU(inplace=True),
-            Conv2dBN(stem_channels, out_channels, 1, 1, 0),
+            Conv2dGN(stem_channels, out_channels, 1, 1, 0),
             nn.ReLU(inplace=True),
-            Conv2dBN(out_channels, out_channels, 3, 2, 1),
+            Conv2dGN(out_channels, out_channels, 3, 2, 1),
             nn.ReLU(inplace=True),
         )
 
@@ -36,7 +36,7 @@ class XNetDownsample2d(nn.Module):
         if mode != "conv":
             raise ValueError(f"Unsupported downsample mode: {mode}")
         self.block = nn.Sequential(
-            Conv2dBN(in_channels, out_channels, 3, 2, 1),
+            Conv2dGN(in_channels, out_channels, 3, 2, 1),
             nn.ReLU(inplace=True),
         )
 
@@ -49,14 +49,14 @@ class XLocalBranch2d(nn.Module):
     def __init__(self, channels: int) -> None:
         super().__init__()
         self.branch3 = nn.Sequential(
-            Conv2dBN(channels, channels, 3, 1, 1, groups=channels),
+            Conv2dGN(channels, channels, 3, 1, 1, groups=channels),
             nn.ReLU(inplace=True),
-            Conv2dBN(channels, channels, 1, 1, 0),
+            Conv2dGN(channels, channels, 1, 1, 0),
         )
         self.branch5 = nn.Sequential(
-            Conv2dBN(channels, channels, 5, 1, 2, groups=channels),
+            Conv2dGN(channels, channels, 5, 1, 2, groups=channels),
             nn.ReLU(inplace=True),
-            Conv2dBN(channels, channels, 1, 1, 0),
+            Conv2dGN(channels, channels, 1, 1, 0),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -114,16 +114,16 @@ class XWaveletBranch2d(nn.Module):
             channels, wavelet_type=wavelet_type, wavelet_level=wavelet_level
         )
         self.ll_proj = nn.Sequential(
-            Conv2dBN(channels, channels, 3, 1, 1),
+            Conv2dGN(channels, channels, 3, 1, 1),
             nn.ReLU(inplace=True),
         )
         self.high_proj = nn.Sequential(
-            Conv2dBN(channels * 3, channels * 3, 3, 1, 1, groups=channels * 3),
+            Conv2dGN(channels * 3, channels * 3, 3, 1, 1, groups=channels * 3),
             nn.ReLU(inplace=True),
-            Conv2dBN(channels * 3, channels * 3, 1, 1, 0),
+            Conv2dGN(channels * 3, channels * 3, 1, 1, 0),
         )
         self.out_proj = nn.Sequential(
-            Conv2dBN(channels, channels, 1, 1, 0),
+            Conv2dGN(channels, channels, 1, 1, 0),
             nn.ReLU(inplace=True),
         )
 
@@ -150,7 +150,7 @@ class XSSMGlobalBranch2d(nn.Module):
         hidden_ratio = max(global_ratio, 1.0)
         self.backend = ssm_backend
         self.pre = nn.Sequential(
-            Conv2dBN(channels, channels, 1, 1, 0),
+            Conv2dGN(channels, channels, 1, 1, 0),
             nn.ReLU(inplace=True),
         )
         self.ssm = VMambaSS2D(
@@ -164,7 +164,7 @@ class XSSMGlobalBranch2d(nn.Module):
             channel_first=True,
         )
         self.post = nn.Sequential(
-            Conv2dBN(channels, channels, 1, 1, 0),
+            Conv2dGN(channels, channels, 1, 1, 0),
             nn.ReLU(inplace=True),
         )
 
@@ -225,7 +225,7 @@ class XBranchFusion2d(nn.Module):
         fused_channels = channels * num_branches
         hidden_channels = max(channels // 4, 8)
         self.fuse = nn.Sequential(
-            Conv2dBN(fused_channels, channels, 1, 1, 0),
+            Conv2dGN(fused_channels, channels, 1, 1, 0),
             nn.ReLU(inplace=True),
         )
         self.gate = nn.Sequential(
@@ -258,7 +258,7 @@ class XTEB2d(nn.Module):
         ssm_backend: str = "auto",
     ) -> None:
         super().__init__()
-        self.pre_norm = Conv2dBN(channels, channels, 1, 1, 0)
+        self.pre_norm = Conv2dGN(channels, channels, 1, 1, 0)
         self.local_branch = XLocalBranch2d(channels)
         self.wavelet_branch = (
             XWaveletBranch2d(
@@ -280,14 +280,14 @@ class XTEB2d(nn.Module):
         )
         self.fusion = XBranchFusion2d(channels, num_branches=3)
         self.post = nn.Sequential(
-            Conv2dBN(channels, channels, 3, 1, 1),
+            Conv2dGN(channels, channels, 3, 1, 1),
             nn.ReLU(inplace=True),
-            Conv2dBN(channels, channels, 1, 1, 0, bn_weight_init=0.0),
+            Conv2dGN(channels, channels, 1, 1, 0, bn_weight_init=0.0),
         )
         self.ffn = nn.Sequential(
-            Conv2dBN(channels, channels * 2, 1, 1, 0),
+            Conv2dGN(channels, channels * 2, 1, 1, 0),
             nn.ReLU(inplace=True),
-            Conv2dBN(channels * 2, channels, 1, 1, 0, bn_weight_init=0.0),
+            Conv2dGN(channels * 2, channels, 1, 1, 0, bn_weight_init=0.0),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -426,15 +426,15 @@ class XSkipFusion2d(nn.Module):
     def __init__(self, in_channels: int, skip_channels: int, out_channels: int) -> None:
         super().__init__()
         self.input_proj = nn.Sequential(
-            Conv2dBN(in_channels, out_channels, 1, 1, 0),
+            Conv2dGN(in_channels, out_channels, 1, 1, 0),
             nn.ReLU(inplace=True),
         )
         self.skip_proj = nn.Sequential(
-            Conv2dBN(skip_channels, out_channels, 1, 1, 0),
+            Conv2dGN(skip_channels, out_channels, 1, 1, 0),
             nn.ReLU(inplace=True),
         )
         self.fuse = nn.Sequential(
-            Conv2dBN(out_channels * 2, out_channels, 3, 1, 1),
+            Conv2dGN(out_channels * 2, out_channels, 3, 1, 1),
             nn.ReLU(inplace=True),
         )
 
@@ -467,9 +467,9 @@ class XFrequencyRefine2d(nn.Module):
             nn.Sigmoid(),
         )
         self.refine = nn.Sequential(
-            Conv2dBN(channels, channels, 3, 1, 1, groups=channels),
+            Conv2dGN(channels, channels, 3, 1, 1, groups=channels),
             nn.ReLU(inplace=True),
-            Conv2dBN(channels, channels, 1, 1, 0),
+            Conv2dGN(channels, channels, 1, 1, 0),
         )
         self.learnable_low_freq_radius = learnable_low_freq_radius
         if learnable_low_freq_radius:
@@ -565,9 +565,9 @@ class XCRB2d(nn.Module):
             else nn.Identity()
         )
         self.out_refine = nn.Sequential(
-            Conv2dBN(out_channels, out_channels, 3, 1, 1),
+            Conv2dGN(out_channels, out_channels, 3, 1, 1),
             nn.ReLU(inplace=True),
-            Conv2dBN(out_channels, out_channels, 3, 1, 1, bn_weight_init=0.0),
+            Conv2dGN(out_channels, out_channels, 3, 1, 1, bn_weight_init=0.0),
         )
 
     def forward(
@@ -586,9 +586,9 @@ class XNetHeadRefine2d(nn.Module):
         if out_channels is None:
             out_channels = channels
         self.block = nn.Sequential(
-            Conv2dBN(channels, out_channels, 3, 1, 1),
+            Conv2dGN(channels, out_channels, 3, 1, 1),
             nn.ReLU(inplace=True),
-            Conv2dBN(out_channels, out_channels, 3, 1, 1),
+            Conv2dGN(out_channels, out_channels, 3, 1, 1),
             nn.ReLU(inplace=True),
         )
 
@@ -662,7 +662,7 @@ class XNetSegHead2d(nn.Module):
     ) -> None:
         super().__init__()
         self.block = nn.Sequential(
-            Conv2dBN(in_channels, in_channels, 3, 1, 1),
+            Conv2dGN(in_channels, in_channels, 3, 1, 1),
             nn.ReLU(inplace=True),
             nn.Conv2d(in_channels, num_classes, kernel_size=1, bias=True),
         )

tests/test_xnet_2d.py

@@ -87,3 +87,58 @@ def test_xnet2d_decoder_uses_plain_unet_skip_connections() -> None:
     decoder_module_names = dict(model.decoder.named_modules())
 
     assert not any(name.startswith("guide") for name in decoder_module_names)
+
+
+def test_xnet2d_main_path_does_not_use_batchnorm2d() -> None:
+    from lib.modules.xnet_2d import XNet2d
+
+    model = XNet2d(
+        in_channels=3,
+        num_classes=1,
+        encoder_channels=(8, 16, 24, 32),
+        encoder_depths=(1, 1, 1, 1),
+        decoder_channels=(24, 16, 8),
+        stem_channels=8,
+        bottleneck_depth=1,
+        use_global_branch_stage1=False,
+        ssm_d_state=1,
+        ssm_backend="torch",
+    )
+
+    batchnorm_modules = [
+        module
+        for module in model.modules()
+        if isinstance(module, nn.BatchNorm2d)
+    ]
+
+    assert not batchnorm_modules
+
+
+def test_xnet2d_train_mode_supports_batch_size_one_without_batchnorm_failure() -> None:
+    from lib.modules.xnet_2d import XNet2d, XTEB2d
+
+    model = XNet2d(
+        in_channels=3,
+        num_classes=1,
+        encoder_channels=(8, 16, 24, 32),
+        encoder_depths=(1, 1, 1, 1),
+        decoder_channels=(24, 16, 8),
+        stem_channels=8,
+        bottleneck_depth=1,
+        global_ratio=1.0,
+        use_wavelet_branch=True,
+        use_global_branch_stage1=False,
+        ssm_d_state=1,
+        ssm_backend="torch",
+        use_frequency_refine=True,
+        learnable_low_freq_radius=False,
+    )
+    for module in model.modules():
+        if isinstance(module, XTEB2d):
+            module.global_branch = nn.Identity()
+    model.train()
+
+    x = torch.randn(1, 3, 64, 64)
+    outputs = model(x)
+
+    assert outputs["seg_logits"].shape == (1, 1, 64, 64)