以下方案为个人试错后经验总结,在下面部署流程中可能遇到若干问题可与主页另一篇《Windows系统本地部署Vision_Mamba问题集锦》配合使用,祝好!
一、检查 Windows 版本
按Win + R,输入:
winver建议 Windows 10 至少是21H2 或更新版本。Ubuntu 官方关于 WSL GPU 加速的说明也把 Windows 10 21H2 作为最低推荐条件之一。
然后用 PowerShell 检查显卡:
nvidia-smi能看到 RTX3050、驱动版本、CUDA Version,就说明 Windows 侧显卡驱动基本可用。注意这里显示的CUDA Version 12.x不代表你必须装 CUDA 12.x,它表示驱动最高支持的 CUDA 版本;我们后面会在环境里使用 PyTorch 官方的cu118版本。
二、安装 WSL2 + Ubuntu
用管理员身份打开 PowerShell,执行:
wsl --install -d Ubuntu-22.04 wsl --update wsl --set-default-version 2 wsl -l -v如果最后看到 Ubuntu 的版本是2,就说明是 WSL2。Microsoft 官方文档也给出的核心安装命令就是wsl --install,并建议用wsl --update更新 WSL 内核。
第一次打开 Ubuntu 时,它会让你创建 Linux 用户名和密码,正常设置即可。
三、在 WSL2 Ubuntu 里安装 CUDA 编译环境
打开 Ubuntu 终端,先确认 WSL 能看到显卡:
nvidia-smi如果能看到 RTX3050,继续。
然后安装基础编译工具:
sudo apt update sudo apt install -y git wget curl build-essential python3.10 python3.10-venv python3.10-dev python3-pip rsync接着安装 CUDA 11.8 Toolkit。注意,这里装的是WSL 专用 CUDA Toolkit,不是 Linux 显卡驱动。NVIDIA 文档明确说,WSL2 里不要安装 Linux display driver,编译 CUDA 程序才需要 CUDA Toolkit。
cd ~ wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600 wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-wsl-ubuntu-11-8-local_11.8.0-1_amd64.deb sudo dpkg -i cuda-repo-wsl-ubuntu-11-8-local_11.8.0-1_amd64.deb sudo cp /var/cuda-repo-wsl-ubuntu-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/ sudo apt-get update sudo apt-get -y install cuda-toolkit-11-8配置环境变量:
echo 'export CUDA_HOME=/usr/local/cuda-11.8' >> ~/.bashrc echo 'export PATH=$CUDA_HOME/bin:$PATH' >> ~/.bashrc echo 'export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc source ~/.bashrc验证:
nvcc -V如果能看到release 11.8,说明 CUDA 编译器准备好了。
四、从零创建 Vision Mamba 项目
建议把项目放在 WSL 的 Linux 文件系统里,不要放在/mnt/c/或/mnt/d/,这样编译 CUDA 扩展更稳定、速度也更好。
mkdir -p ~/projects cd ~/projects git clone https://github.com/hustvl/Vim.git cd Vim最终项目目录大概是:
~/projects/Vim ├── causal-conv1d ├── mamba-1p1p1 ├── vim ├── assets ├── det ├── seg └── ...Vision Mamba 官方仓库的 README 里写明了 NVIDIA GPU 环境建议使用 Python 3.10.13、PyTorch 2.1.1 + cu118,并要求安装causal_conv1d和mamba-1p1p1。
五、创建 Python 虚拟环境
在~/projects/Vim目录下执行:
python3.10 -m venv .venv source .venv/bin/activate升级基础工具:
pip install --upgrade pip pip install setuptools==69.5.1 wheel packaging ninja这里固定setuptools==69.5.1是为了减少老 CUDA 扩展编译时出现pkg_resources或setup.py install is deprecated相关问题。你之前遇到过No module named pkg_resources,本质就是 setuptools 相关环境不完整。
六、安装 PyTorch 2.1.1 + CUDA 11.8
执行:
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118这个版本和 Vision Mamba 官方 README 里的推荐版本一致。
验证 PyTorch 是否调用 GPU:
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"正常输出应该类似:
2.1.1+cu118 True NVIDIA GeForce RTX 3050 ...七、安装 Vision Mamba 的核心 CUDA 扩展
这是最容易失败的地方。一定要在 WSL Ubuntu 里编译,不要在 Windows 原生 PowerShell 里编译。
先设置 RTX3050 的 CUDA 架构,减少编译压力:
export TORCH_CUDA_ARCH_LIST="8.6"然后安装causal-conv1d:
cd ~/projects/Vim/causal-conv1d pip install -e . --no-build-isolation再安装mamba-1p1p1:
cd ~/projects/Vim/mamba-1p1p1 pip install -e . --no-build-isolation--no-build-isolation很重要。Mamba 官方 README 也说明这个参数的作用是让 pip 使用当前环境里已经安装好的 CUDA 版 PyTorch,而不是在隔离环境里错误地拉 CPU 版 torch。
八、安装 Python 依赖
先回到项目根目录:
cd ~/projects/Vim source .venv/bin/activate为了先把分类模型跑通,建议先装最小必要依赖:
pip install timm==0.4.12 einops==0.7.0 transformers==4.35.2 datasets==2.15.0 mlflow==2.9.1 opencv-python==4.8.1.78 pillow==10.1.0 scipy==1.11.4 scikit-learn==1.3.2 pandas==2.1.3 numpy==1.26.2 pyyaml==6.0.1 tqdm==4.66.1官方也提供了vim/vim_requirements.txt,但它包含较多检测、分割、日志相关依赖,例如mmcv、mmsegmentation等。你现在如果只是先部署并跑通 Vim 分类模型,不建议一开始就把所有重依赖全装上;等分类模型确认能跑,再补装检测/分割部分更稳。
九、写一个最小测试脚本
在项目根目录创建文件:
cd ~/projects/Vim nano quick_test_vim.py粘贴下面内容:
import os import sys import torch # 让 Python 能找到 vim/models_mamba.py 和 vim/rope.py ROOT = os.path.dirname(os.path.abspath(__file__)) sys.path.append(os.path.join(ROOT, "vim")) import models_mamba # noqa: F401 from timm.models import create_model def main(): print("Torch:", torch.__version__) print("CUDA available:", torch.cuda.is_available()) print("GPU:", torch.cuda.get_device_name(0)) model_name = "vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2" model = create_model( model_name, pretrained=False, num_classes=1000, ).cuda().eval() x = torch.randn(1, 3, 224, 224).cuda() with torch.no_grad(): y = model(x) print("Output shape:", y.shape) if __name__ == "__main__": main()保存后运行:
python quick_test_vim.py如果最后输出:
Output shape: torch.Size([1, 1000])说明 Vision Mamba 的核心模型、Mamba CUDA 扩展、PyTorch GPU 环境已经跑通了。
十、下载官方预训练权重
Vision Mamba 官方 Hugging Face 页面提供了 Vim-tiny、Vim-small、Vim-base 等权重。比如 Vim-tiny 是 ImageNet-1K 上的通用视觉 backbone,Hugging Face 模型卡说明它用于分类、检测、分割等视觉任务。
创建权重目录:
cd ~/projects/Vim mkdir -p checkpoints下载 Vim-tiny:
wget -O checkpoints/vim_t_midclstok_76p1acc.pth "https://huggingface.co/hustvl/Vim-tiny-midclstok/resolve/main/vim_t_midclstok_76p1acc.pth?download=true"如果你想下微调后的 tiny 权重:
wget -O checkpoints/vim_t_midclstok_ft_78p3acc.pth "https://huggingface.co/hustvl/Vim-tiny-midclstok/resolve/main/vim_t_midclstok_ft_78p3acc.pth?download=true"Vim-small 权重大约 400MB 多,RTX3050 也能推理,但训练会更吃显存。Hugging Face 页面显示 Vim-small 有vim_s_midclstok_80p5acc.pth和vim_s_midclstok_ft_81p6acc.pth。
十一、测试加载预训练权重
创建:
nano quick_test_vim_pretrained.py粘贴:
import os import sys import torch ROOT = os.path.dirname(os.path.abspath(__file__)) sys.path.append(os.path.join(ROOT, "vim")) import models_mamba # noqa: F401 from timm.models import create_model def main(): model_name = "vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2" ckpt_path = os.path.join(ROOT, "checkpoints", "vim_t_midclstok_76p1acc.pth") model = create_model( model_name, pretrained=False, num_classes=1000, ) checkpoint = torch.load(ckpt_path, map_location="cpu") state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint missing, unexpected = model.load_state_dict(state_dict, strict=False) print("Missing keys:", len(missing)) print("Unexpected keys:", len(unexpected)) model = model.cuda().eval() x = torch.randn(1, 3, 224, 224).cuda() with torch.no_grad(): y = model(x) print("Output shape:", y.shape) print("Top class index:", y.argmax(dim=1).item()) if __name__ == "__main__": main()运行:
python quick_test_vim_pretrained.py能输出torch.Size([1, 1000]),就说明预训练权重也加载成功了。
十二、把项目接入 PyCharm
PyCharm 官方文档说明,可以通过On WSL添加 WSL 里的 Python 解释器。
推荐做法:
打开 PyCharm,选择:
Open然后打开这个路径:
\\wsl$\Ubuntu-22.04\home\你的WSL用户名\projects\Vim接着配置解释器:
Settings → Project: Vim → Python Interpreter → Add Interpreter → On WSL → 选择 Ubuntu-22.04 → Existing environment → /home/你的WSL用户名/projects/Vim/.venv/bin/python如果你的 PyCharm 里没有On WSL,通常是版本或授权功能限制。JetBrains 文档写明 PyCharm Pro 支持通过 WSL、SSH、Docker 等方式配置远程/WSL 解释器。
PyCharm 里运行测试脚本时,配置如下:
Script path: ~/projects/Vim/quick_test_vim.py Working directory: ~/projects/Vim Python interpreter: /home/你的WSL用户名/projects/Vim/.venv/bin/python如果 PyCharm 的路径选择界面显示 Windows 路径,不要选D:\PythonProject\Vim\.venv\Scripts\python.exe,要选 WSL 里的:
/home/你的WSL用户名/projects/Vim/.venv/bin/python十三、跑官方 main.py 做验证
Vision Mamba 官方 README 给的评估命令形式是:
python main.py --eval --resume /path/to/ckpt --model vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2 --data-path /path/to/imagenet官方就是用这个方式评估 ImageNet-1K 权重。
如果你没有 ImageNet,就先不要跑这个命令,因为它需要完整 ImageNet 验证集目录。你现在第一目标应该是先用quick_test_vim.py跑通模型结构和 GPU。
如果以后你要用 CIFAR10 训练,可以再改官方main.py参数。RTX3050 显存有限,建议先从很小的 batch 开始,例如:
cd ~/projects/Vim/vim source ../.venv/bin/activate python main.py \ --model vim_tiny_patch16_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2 \ --data-set CIFAR \ --data-path ./data \ --batch-size 4 \ --epochs 1 \ --num_workers 2 \ --output_dir ../output/test_cifar \ --no_amp如果显存不够,把--batch-size 4改成2或1。