UMamba in encoder

[anwai98]

UMambaEnc is a little weird. The training didn't work at all for 3/5 folds.

Suspicions:

• Unbalanced train -> train-val split, as I let nnUNet's random splitting framework take care of it for the 5-fold cross validation.
  • AA: I don't think this is the case, because it does the same for UMambaBot, and all the trainings worked there just fine. If random splitting was the case, I must have seen it in atleast either of the UMambaBot training folds.
• Would try to infer now on labeled-test set, and see what's going on (maybe report and discuss)

PS: UMambaEnc - mamba implementation in the entire encoder; UMambaBot - mamba bottleneck layer b/w encoder and decoder

[anwai98]

Other users have faced the issue as well:

1. Ref: UMamba - Nan When training bowang-lab/U-Mamba#8
2. Ref: Vim - Loss is nan, stopping training hustvl/Vim#30

Mention: I have faced Point 1 extensively, and Point 2 in ViMUNet for the largest vim model (vim_b)

Edit: Both the issues are taken care of now (in U-Mamba and ViM-UNet)

[anwai98]

Also, a nice spot to track Vision Mamba-related work: https://github.com/VisionMamba