Vision-RWKV

“Vision-RWKV：基于RWKV类似架构的高效可扩展视觉感知”（arXiv:2403.02308）的官方实现。

新闻🚀🚀🚀

• 2025/02/18：在cuda_new文件夹中新增了CUDA代码版本，以消除T_MAX的硬编码。
• 2025/02/11：🎊🎊 Vision-RWKV已被ICLR 2025接收！
• 2024/04/14：我们在分类任务中支持rwkv6，性能更高！
• 2024/03/04：我们发布了Vision-RWKV的代码和模型。

亮点

• 高分辨率效率：以全局感受野流畅处理高分辨率图像。
• 可扩展性：基于大规模数据集预训练，并具备良好的规模扩展稳定性。
• 卓越性能：在分类任务中表现优于ViT；在密集预测任务中，其FLOPs更低、速度更快，超越了基于窗口的ViT，与全局注意力ViT相当。
• 高效替代方案：可在综合性视觉任务中作为ViT的替代骨干网络。

概述

进度计划

• 支持RWKV6作为VRWKV6
• 发布VRWKV-L
• 发布VRWKV-T/S/B

模型库

预训练模型

模型	尺寸	预训练数据集	下载链接
VRWKV-L	192	ImageNet-22K	ckpt

图像分类（ImageNet-1K）

模型	尺寸	参数量	FLOPs	Top-1准确率	下载链接
VRWKV-T	224	6.2M	1.2G	75.1	ckpt | 配置文件
VRWKV-S	224	23.8M	4.6G	80.1	ckpt | 配置文件
VRWKV-B	224	93.7M	18.2G	82.0	ckpt | 配置文件
VRWKV-L	384	334.9M	189.5G	86.0	ckpt | 配置文件
VRWKV6-T	224	7.6M	1.6G	76.6	ckpt | 配置文件
VRWKV6-S	224	27.7M	5.6G	81.1	ckpt | 配置文件
VRWKV6-B	224	104.9M	20.9G	82.6	ckpt | 配置文件

• VRWKV-L先在ImageNet-22K上预训练，随后在ImageNet-1K上进行微调。
• 我们使用internimage代码库训练VRWKV-L，以提升速度。

目标检测（带Mask-RCNN头，COCO数据集）

模型	参数量	FLOPs	box AP	mask AP	下载链接
VRWKV-T	8.4M	67.9G	41.7	38.0	ckpt | 配置文件
VRWKV-S	29.3M	189.9G	44.8	40.2	ckpt | 配置文件
VRWKV-B	106.6M	599.0G	46.8	41.7	ckpt | 配置文件
VRWKV-L	351.9M	1730.6G	50.6	44.9	ckpt | 配置文件

• 本表报告的是骨干网络的参数量和FLOPs。

语义分割（带UperNet头，ADE20K数据集）

模型	参数量	FLOPs	mIoU	下载链接
VRWKV-T	8.4M	16.6G	43.3	ckpt | 配置文件
VRWKV-S	29.3M	46.3G	47.2	ckpt | 配置文件
VRWKV-B	106.6M	146.0G	49.2	ckpt | 配置文件
VRWKV-L	351.9M	421.9G	53.5	ckpt | 配置文件

• 本表报告的是骨干网络的参数量和FLOPs。

引用

如果本工作对您的研究有所帮助，请考虑引用以下BibTeX条目。

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

许可证

本仓库采用Apache 2.0许可证发布，详情请参阅LICENSE文件。

致谢

Vision-RWKV 的构建参考了以下项目的代码：RWKV、MMPretrain、MMDetection、MMSegmentation、ViT-Adapter、InternImage。感谢这些项目团队的出色工作！

Vision-RWKV 快速上手指南

Vision-RWKV 是一种基于 RWKV 架构的高效视觉感知模型，具备高分辨率处理能力、优秀的扩展性以及在分类、检测和分割任务中超越传统 ViT 的性能。本指南将帮助您快速部署并使用该工具。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• Python: 3.8 或更高版本
• GPU: NVIDIA GPU (支持 CUDA)，显存建议 8GB 以上（根据模型大小而定）
• 核心依赖:
  • PyTorch (建议 1.12+)
  • CUDA Toolkit (需与 PyTorch 版本匹配)
  • MMPretrain, MMDetection, MMSegmentation (用于下游任务)

提示：本项目依赖特定的 CUDA 算子。如果您在中国大陆，建议在安装 Python 依赖时使用国内镜像源（如清华源或阿里源）以加速下载。

安装步骤

1. 克隆项目代码

git clone https://github.com/OpenGVLab/Vision-RWKV.git
cd Vision-RWKV

2. 创建虚拟环境并安装基础依赖

conda create -n vrwkv python=3.9 -y
conda activate vrwkv

# 安装 PyTorch (请根据您的 CUDA 版本选择对应的安装命令，此处以 CUDA 11.8 为例)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装其他 Python 依赖 (推荐使用国内镜像加速)
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 编译 CUDA 算子

Vision-RWKV 需要编译自定义 CUDA 内核以支持高效推理。

# 进入 cuda 目录并编译
cd cuda
python setup.py install
cd ..

# 如果使用最新的 cuda_new 版本 (消除 T_MAX 硬编码限制)
cd cuda_new
python setup.py install
cd ..

4. 安装下游任务框架 (可选)

如果您计划进行目标检测或语义分割，需安装 OpenMMLab 系列库：

# 安装 MMPretrain (分类)
pip install mmpretrain -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 MMDetection (检测) 和 MMSegmentation (分割)
# 请参考各官方仓库的安装指南，或使用项目提供的 requirements 文件
pip install -U openmim
mim install mmdet
mim install mmsegmentation

基本使用

以下示例展示如何加载预训练模型并进行图像分类推理。

1. 下载预训练模型

从 Model Zoo 中选择适合您任务的模型，例如 VRWKV-T (Tiny 版本)：

mkdir checkpoints
wget https://huggingface.co/OpenGVLab/Vision-RWKV/resolve/main/vrwkv_t_in1k_224.pth -O checkpoints/vrwkv_t_in1k_224.pth

(注：如果 HuggingFace 下载缓慢，可尝试使用国内镜像站或代理)

2. Python 推理示例

创建一个 demo.py 文件，运行以下代码：

import torch
from PIL import Image
from torchvision import transforms
from classification.models.vrwkv import VRWKV

# 配置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化模型 (以 VRWKV-Tiny 为例)
# 具体参数需参考对应的配置文件 (.py)
model = VRWKV(
    img_size=224,
    patch_size=16,
    embed_dims=192,
    depth=12,
    num_heads=6,
    # ... 其他参数请参考源码或配置文件
)

# 加载权重
checkpoint = torch.load('checkpoints/vrwkv_t_in1k_224.pth', map_location=device)
model.load_state_dict(checkpoint['state_dict'], strict=False)
model.to(device)
model.eval()

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载图片
image = Image.open("your_image.jpg").convert("RGB")
input_tensor = transform(image).unsqueeze(0).to(device)

# 推理
with torch.no_grad():
    output = model(input_tensor)
    prediction = torch.argmax(output, dim=1)

print(f"预测类别索引: {prediction.item()}")

3. 使用配置文件运行 (推荐)

对于标准的分类、检测和分割任务，建议使用项目提供的配置文件和训练/测试脚本，以确保参数对齐。

图像分类测试示例：

python classification/tools/test.py \
    classification/configs/vrwkv/vrwkv_tiny_8xb128_in1k.py \
    checkpoints/vrwkv_t_in1k_224.pth \
    --eval accuracy

目标检测测试示例 (Mask R-CNN):

python detection/tools/test.py \
    detection/configs/mask_rcnn/mask_rcnn_vrwkv_adapter_tiny_fpn_1x_coco.py \
    checkpoints/mask_rcnn_vrwkv_adapter_tiny_fpn_1x_coco.pth \
    --eval bbox segm