Awesome
MLLA-UNet
The codes for the work "MLLA-UNet: Mamba-like Linear Attention in an Efficient U-Shape Model for Medical Image Segmentation".
1. Pretrained Models: You can choose a pretrained model based on your preference.
| model | Resolution | #Params | FLOPs | acc@1 | config | pretrained weights |
|---|---|---|---|---|---|---|
| MLLA-T | 224 | 25M | 4.2G | 83.5 | config | TsinghuaCloud |
| MLLA-S | 224 | 43M | 7.3G | 84.4 | config | TsinghuaCloud |
| MLLA-B | 224 | 96M | 16.2G | 85.3 | config | TsinghuaCloud |
Ref: [MLLA Official Implementation]
2. Prepare data
Preprocessing:
Ref: nnUNet
3. Environment
- Please prepare an environment with python=3.9 and then use the command
pip install -r requirements.txtfor the dependencies.
4. Train/Test
-
Run the train script on synapse dataset. The batch size we used is 48. If you do not have enough GPU memory, the bacth size can be reduced to 12 or 6 to save memory.
-
Train
sh train.sh
- Test
sh test.sh