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

• AMOS 22
• FLARE 22
• ATLAS
• WORD
• BTCV
• ACDC

Preprocessing:

Ref: nnUNet

3. Environment

• Please prepare an environment with python=3.9 and then use the command pip install -r requirements.txt for 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

References

• SwinUnet
• MLLA