Vision-RWKV

The official implementation of "Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures".

News🚀🚀🚀

• 2025/02/18: A new version of the CUDA code has been added in the cuda_new folder to eliminate the hardcoding of T_MAX.
• 2025/02/11: 🎊🎊 Vison-RWKV is accepted by ICLR 2025!
• 2024/04/14: We support rwkv6 in classification task, higher performance!
• 2024/03/04: We release the code and models of Vision-RWKV.

Highlights

• High-Resolution Efficiency: Processed high-resolution images smoothly with a global receptive field.
• Scalability: Pre-trained with large-scale datasets and posses scale up stablity.
• Superior Performance: Achieved a better performance in classfication tasks than ViTs. Surpassed window-based ViTs and comparabled to global attention ViTs with lower flops and higher speed in dense prediction tasks.
• Efficient Alternative: Capability to be an alternative backbone to ViT in comprehensive vision tasks.

Overview

Schedule

• Support RWKV6 as VRWKV6
• Release VRWKV-L
• Release VRWKV-T/S/B

Model Zoo

Pretrained Models

Image Classification (ImageNet-1K)

• VRWKV-L is pretrained on ImageNet-22K and then finetuned on ImageNet-1K.
• We train VRWKV-L with the internimage codebase for a higher speed.

Object Detection with Mask-RCNN head (COCO)

• We report the #Param and #FLOPs of the backbone in this table.

Semantic Segmentation with UperNet head (ADE20K)

• We report the #Param and #FLOPs of the backbone in this table.

Citation

If this work is helpful for your research, please consider citing the following BibTeX entry.

@article{duan2024vrwkv,
  title={Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures},
  author={Duan, Yuchen and Wang, Weiyun and Chen, Zhe and Zhu, Xizhou and Lu, Lewei and Lu, Tong and Qiao, Yu and Li, Hongsheng and Dai, Jifeng and Wang, Wenhai},
  journal={arXiv preprint arXiv:2403.02308},
  year={2024}
}

License

This repository is released under the Apache 2.0 license as found in the LICENSE file.

Acknowledgement

Vision-RWKV is built with reference to the code of the following projects: RWKV, MMPretrain, MMDetection, MMSegmentation, ViT-Adapter, InternImage. Thanks for their awesome work!