pytorch-image-models
pytorch-image-models copied to clipboard
[FEATURE] Support for RWKV and MAMBA architecture
Is your feature request related to a problem? Please describe. RWKV and MAMBA are new families of models that are efficient. It would be great if TIMM lib have this
Describe the solution you'd like add a new file in timm to implement them.
Describe alternatives you've considered
Additional context Here is the official lib of Vision-RWKV https://github.com/OpenGVLab/Vision-RWKV , and here is the paper https://openreview.net/forum?id=nGiGXLnKhl . And for RWKV this is the website https://rwkv.cn/ .
As for mamba, there are a lot of vision implementations. For example
- https://github.com/hustvl/Vim
- https://github.com/MzeroMiko/VMamba
- https://github.com/NVlabs/MambaVision With papers
- https://icml.cc/virtual/2024/poster/33768
- https://openreview.net/forum?id=ZgtLQQR1K7&referrer=%5Bthe%20profile%20of%20Lingxi%20Xie%5D(%2Fprofile%3Fid%3D~Lingxi_Xie1)
- https://arxiv.org/abs/2407.08083
@2catycm they are both interesting model families, the problem is they all require custom kernels (or external libraries with custom kernels). Those have proven to be difficult to maintain over time, other models requiring custom kernels have come and gone in the time frame timm has existed and original authors did not maintain those kernels for current iterations of cuda/torch.