MrBlankness / LightM-UNet

Pytorch implementation of "LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation"
https://arxiv.org/abs/2403.05246
Apache License 2.0
299 stars 26 forks source link

NnUNetTrainerLightMUNet has not been used #32

Open AFuJianPeople opened 3 months ago

AFuJianPeople commented 3 months ago

Hello, I'm very sorry to disturb you. I didn't see nnUNetTrainerLightMUNet being used in the code, it has always been nnUNetTrainer. Could you tell me where nnUNetTrainerLightMUNet is being used? Thank you very much.

alvinkimbowa commented 3 months ago

It is passed as an argument during training as mentioned in the README:

nnUNetv2_train DATASET_ID 2d all -tr nnUNetTrainerLightMUNet

By default nnUNetv2_train uses nnUNetTrainer, and if you want to incorporate a custom model such as LightM-UNet, a quick and dirty way is to write a custom trainer that will train using your custom model. All other stages in the pipeline (data preprocessing and augmentation) remain the same, except the model training. Please take a look at the nnUNet documentation - Extending nnUNet.

AFuJianPeople commented 3 months ago

It is passed as an argument during training as mentioned in the README:

nnUNetv2_train DATASET_ID 2d all -tr nnUNetTrainerLightMUNet

By default nnUNetv2_train uses nnUNetTrainer, and if you want to incorporate a custom model such as LightM-UNet, a quick and dirty way is to write a custom trainer that will train using your custom model. All other stages in the pipeline (data preprocessing and augmentation) remain the same, except the model training. Please take a look at the nnUNet documentation - Extending nnUNet.

Thank you very much for your answer. With your help, I have started modifying my own model.