X_SSL_Net/docs/archive/20260504_220523_xnet_vmamba_cleanup_archive.md

XNet + VMamba 接入与清理归档（2026-05-04）

1. 本轮目标

本轮工作的核心目标不是继续发散设计，而是把 XNet2d 从概念结构推进到一版可运行、可训练、来源清晰、依赖边界尽量干净的工程原型。

重点围绕四件事展开：

• 明确 XNet 的方法学结构；
• 实现 XNet2d 初版；
• 将 global branch 从临时轻量 mixer 升级到真正的 VMamba SS2D；
• 清理 lib/modules 与 lib_mamba，使当前工程聚焦到 XNet 主路径。

2. 架构结论

2.1 XNet 的结构定义已经基本定型

当前 XNet 的核心结构共识如下：

• 编码器 block 同构：统一使用 X-TEB
• 解码器 block 同构：统一使用 X-CRB
• 编码器和解码器 block 不同：属于非对称 encoder-decoder
• 仍保留 U-Net 的层级和 skip 稳定性
• 但 decoder 不只接收同尺度 skip，还接收更深层语义 guide
• X 的来源是：
  • E_i -> D_i 的同尺度路径
  • E_{i+1} -> D_i 的斜向语义引导路径

这意味着它不是改名 U-Net，而是一个带显式 cross-scale guide 的 X-shaped 编解码结构。

2.2 分支职责已经明确

X-TEB 当前角色分工：

• local branch：多核 depthwise 局部建模
• wavelet branch：结构性频率建模
• global branch：VMamba SS2D

X-CRB 当前角色分工：

• 上采样恢复
• same-scale skip 融合
• diagonal guide 调制
• 频率细化

3. 文档状态

已在：

tmp/docs/method/XNet_method.md

中形成一版较完整的方法草稿，包含：

• 中文方法定义
• 张量尺寸流
• 纯文本结构图
• X-TEB / X-CRB 公式化说明
• PyTorch 接口草稿
• 训练与消融建议

另补充了来源说明：

tmp/docs/method/lib_mamba_source_note.md

其作用是记录：

• lib_mamba 当前主接入来源
• VMamba 与 MobileMamba 的关系边界
• 当前 XNet2d 使用的是哪一条 SS2D 分支

4. 代码实现状态

4.1 XNet2d 已实现

主实现位于：

lib/modules/xnet_2d.py

当前已实现的主要类：

• XNetStem2d
• XNetDownsample2d
• XWaveletBranch2d
• XSSMGlobalBranch2d
• XGlobalBranch2d
• XTEB2d
• XNetEncoderStage2d
• XNetEncoder2d
• XGuideProjector2d
• XCRB2d
• XNetDecoder2d
• XNetSegHead2d
• XNetBoundaryHead2d
• XNet2d

4.2 当前 global branch 已切到正统 VMamba 路线

这是本轮最关键的工程结论之一。

最初做过一个临时轻量 token mixer 版 global branch，但后续已经清理掉。

当前 XNet2d 的 global branch 已接入：

tmp/VMamba-main/classification/models/vmamba.py

对应的 SS2D 主线实现，而不是：

• MobileMamba 的 vmambanew.py
• 也不是其它再改写版本

这使得方法归因更加清晰：

• XNet 的 global branch 是基于 VMamba 主线 SS2D
• 而不是基于 MobileMamba 的二次改写分支

5. 关于 SS2D 来源的判断结论

用户特别关心的一点已经明确：

5.1 SS2D 不是 Mamba 官方原生模块

当前判断是：

• SS2D 不是 state-spaces/mamba 官方仓库里的原生模块
• 它属于 VMamba 这条视觉 2D 状态空间扩展路线
• 底层 selective_scan 内核与 Mamba 官方体系相关
• 但 SS2D 本身是视觉化、二维化后的实现，不应表述成“官方 Mamba 2D 实现”

5.2 当前最合理表述

后续写文档/论文时，应表述为：

• global branch adopts a VMamba-style SS2D operator

而不要写成：

• official Mamba 2D implementation

6. lib_mamba 当前状态

6.1 已做的接入

当前项目内已有：

lib/modules/lib_mamba

并且主接入分支已经切换到：

• vmamba.py
• csm_triton.py
• csms6s.py

6.2 已做的最小清理

本轮已将 lib/modules/lib_mamba/__init__.py 简化为：

• 只导出 SS2D

同时按 XNet 当前需要做过一次删减和收缩。

但需要注意：

• 用户后续又自行恢复过 kernels/selective_scan
• 因此当前 lib_mamba 里仍存在编译后的 selective_scan 相关目录
• 这不是错误，而是用户为了使用 oflex 主动恢复的合理内容

7. selective_scan_cuda_oflex 结论

7.1 是否已经可用

已经确认：

• selective_scan_cuda_oflex 可以在正确加载 torch 后成功导入
• csms6s.WITH_SELECTIVESCAN_OFLEX == True

用户验证结果表明：

python -c "import torch; import selective_scan_cuda_oflex; print('oflex OK after torch')"

是成功的。

7.2 为什么裸导入一度报 libc10.so

原因不是编译失败，而是：

• 直接 import selective_scan_cuda_oflex
• 但尚未先加载 torch
• 导致 PyTorch 动态库路径未初始化

因此：

• 先 import torch
• 再 import selective_scan_cuda_oflex

即可正常导入。

7.3 当前真正结论

当前 XNet2d -> SS2D -> selective_scan 这条链路已经能够吃到 oflex。

8. 环境验证结论

8.1 torch_py310_cuda126 最终恢复可用

中途一度判断 torch_py310_cuda126 环境已损坏，但用户后续确认该环境在其本机上最终可用：

• torch.__version__ == 2.11.0+cu126
• torch.cuda.is_available() == True
• GPU 为 RTX 4070 Ti SUPER

因此，后续应以用户本机环境验证结果为准，而不是将当前 agent 侧 conda run 的 GPU 不可见误判为环境损坏。

8.2 GPU 前向验证

用户最终已验证：

python -c "import torch; from lib.modules import XNet2d; model = XNet2d(...).cuda().eval(); ..."

输出：

(1, 2, 128, 128)

说明：

• XNet2d GPU 前向可通过
• 输出尺寸正确

9. lib/modules 清理结论

9.1 当前导出边界已明显收缩

lib/modules/__init__.py 目前已改成只导出：

• 基础层
• XNet 相关类

不再导出：

• Swin 编码器
• 旧 segmentation model
• 其它与当前 XNet 无关的模块

9.2 trainer 已切到 XNet2d

lib/trainers/supervised.py

已经从旧的：

• SegmentationModel2d

切到：

• XNet2d

并且构造参数也已调整成读取 XNet 所需字段。

这意味着：

• 当前 lib/modules 里大量旧的 Swin 依赖即使不存在
• 主训练入口也不会再直接引用它们

10. 当前仍存在但可接受的警告

10.1 oflex unavailable 警告

在某些 agent 侧环境中仍可见：

selective_scan_cuda_oflex is unavailable; SS2D will use a slower backend.

但就用户本机验证结果而言：

• oflex 实际可用
• 该 warning 不应再作为主判断依据

10.2 timm 的 warning 已处理

之前存在：

• timm.models.layers 的废弃导入 warning

已改用：

• timm.layers

所以这类 warning 已基本清掉。

10.3 torch.amp 警告已处理

之前 csms6s.py 里使用：

• torch.cuda.amp.custom_fwd/custom_bwd

已改成：

• torch.amp.custom_fwd/custom_bwd

因此这部分警告已处理。

11. 当前代码状态的简洁判断

如果只看对下一轮最重要的结论，可压缩成下面几句：

1. XNet2d 已实现，且 GPU 前向通过。
2. global branch 已切到 VMamba-main/classification/models/vmamba.py 的 SS2D 主线。
3. selective_scan_cuda_oflex 已在用户环境中确认可用。
4. lib/modules 已清成以 XNet 为核心的较小集合。
5. lib/trainers/supervised.py 已改成使用 XNet2d。

12. 当前最合理的下一步

如果下一轮继续做工程推进，优先级建议如下：

12.1 第一优先级

把 XNet2d 接入现有训练配置并实际跑最小训练实验。

这比继续做结构层面的发散更重要，因为：

• 主干实现已经具备；
• GPU 前向已经通；
• SS2D 主路径已经接上；
• 下一步最该验证的是训练稳定性与 loss 收敛。

12.2 第二优先级

进一步裁剪 vmamba.py，只保留：

• SS2Dv2
• SS2D
• 必要辅助函数

删除：

• VSSM
• SS2Dv3
• SS2Dm0
• 与 XNet 无关的工厂函数

这样可以进一步让 lib_mamba 变成 XNet 私有最小版本。

12.3 第三优先级

为 XNet2d 增加一版专用 config 模板，并开始：

• wavelet branch on/off
• frequency refine on/off
• SS2D on/off
• diagonal guide on/off

这四类核心消融。

13. 一句总括

本轮工作的真正精华结论是：

XNet 的结构已经从概念草图推进到工程原型，且 global branch 已被成功替换为基于 VMamba 主线 SS2D 的 GPU 可运行实现；同时项目的模块导入边界与训练入口已开始围绕 XNet 收缩，具备进入训练与消融阶段的条件。