pytorch-image-models icon indicating copy to clipboard operation
pytorch-image-models copied to clipboard

[FEATURE] Support for RWKV and MAMBA architecture

Open 2catycm opened this issue 11 months ago • 1 comments

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 avatar Dec 12 '24 10:12 2catycm

@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.

rwightman avatar Dec 12 '24 17:12 rwightman