（简单易学）NdMamba2的纯PyTorch实现的多维实现。任意模态（signal | nlp | cv | audio| vedio）的快速Mamba2缝合神器（支持1d,2d,3d,...Nd）

提示：文章底部有项目链接

BiMamba2 for any dimension by pytorch

仅使用PyTorch实现的双向Mamba2（BiMamba2）提供N维度支持，包括1d，2d，3d数据的支持，借助BiMamba2你可以很方便的缝合到任意模型中去提高精度。

nd_mamba2.py特点（更新时间：2024/08/02）

• ✅ 支持定制的1d数据(batch,channel,length)
• ✅ 支持定制的2d数据(batch,channel,height,width)
• ✅ 支持定制的3d数据(batch,channel,deep,height,width)
• ✅ 支持通用的Nd数据(batch,channel,*size)
• ✅ 有好的环境支持（纯PyTorch实现，即插即用）

ex_bi_mamba2.py新特性（更新时间：2024/08/16）

• ✅ 支持torch.jit.scipt格式导出（取消了einops库以及配置类）
• ✅ 支持onnx格式导出（采用onnx_14的版本，支持下三角阵的操作）
• ✅ 更易阅读 (删除了大量冗余代码)
• ✅ 兼容nd_mamba2.py的所有特点

bi_mamba2_ac.py新特性（更新时间：2024/08/20）

• ✅ 更好的2d图像语义表达（使用非对称卷积的策略优化2d）
• ✅ 兼容ex_bi_mamba2.py中的所有特点
• ❌ 不支持1d、3d等其他维度的数据格式

vssd_torch.py新特性（更新时间：2024/08/21）

• ✅ 更好的2d图像语义表达（基于VSSD的优化）
• ✅ 友好的环境支持（纯torch实现，即插即用）
• ❌ 不支持导出为scipt及onnx格式
• ❌ 不支持1d、3d等其他维度的数据格式

ex_vssd.py新特性（更新时间：2024/08/22）

• ✅ 更易阅读 (对vssd_torch进行了代码优化)
• ✅ 支持导出为scipt及onnx格式
• ✅ 兼容ex_vssd.py中的所有特点
• ❌ 不支持1d、3d等其他维度的数据格式

提示

如果你想要更快的速度，可以将本项目中的Mamba2替换为Mamba2官方的Cuda加速实现，并按照官方要求安装各种依赖包，这不会影响本项目对多维度数据的支持，但会对模型的导出产生影响

致谢

• [Albert Gu], [Tri Dao] state-spaces/mamba - authors of the Mamba-2 architecture
• [Yuheng Shi], [Minjing Dong] YuHengsss/VSSD - authors of the vssd architecture
• [Thomas] - author of tommyip/mamba2-minimal, who inspired this repo

该项目链接：https://github.com/Human9000/nd-Mamba2-torch