LMa-UNet: First large kernel Mamba for medical segmentation elevates SSMs beyond Convolution and Self-attention 🚀
Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention
python 3.10 + torch 2.0.1 + torchvision 0.15.2 (cuda 11.8)
If cuda is 11.8 run:
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 -f https://download.pytorch.org/whl/torch_stable.html
Install Mamba: pip install causal-conv1d
and pip install mamba-ssm
Install monai: pip install monai
Download code: git clone https://github.com/wjh892521292/LMa-UNet
and cd LMa-UNet/lmaunet
and run pip install -e .
nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity
- Train 2D
LMaUNet
model
nnUNetv2_train DATASET_ID 2d all -tr nnUNetTrainerLMaUNet
- Train 3D
LMaUNet
model
nnUNetv2_train DATASET_ID 3d_fullres all -tr nnUNetTrainerLMaUNet
- Predict testing cases with
LMaUNet
model
nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_ID -c CONFIGURATION -f 'all' -tr nnUNetTrainerLMaUNet --disable_tta -npp 1
CONFIGURATION
can be2d
and3d_fullres
for 2D and 3D models, respectively.
@article{wang2024large,
title={Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention},
author={Jinhong Wang and Jintai Chen and Danny Chen and Jian Wu},
journal={arXiv preprint arXiv:2403.07332},
year={2024}
}
Thank the authors of nnU-Net, Mamba and U-mamba for making their valuable code publicly available.