以下方案为个人试错后经验总结在下面部署流程中可能遇到若干问题可与主页另一篇《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 setuptools69.5.1 wheel packaging ninja这里固定setuptools69.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 torch2.1.1 torchvision0.16.1 torchaudio2.1.1 --index-url https://download.pytorch.org/whl/cu118这个版本和 Vision Mamba 官方 README 里的推荐版本一致。验证 PyTorch 是否调用 GPUpython -c import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))正常输出应该类似2.1.1cu118 True NVIDIA GeForce RTX 3050 ...七、安装 Vision Mamba 的核心 CUDA 扩展这是最容易失败的地方。一定要在 WSL Ubuntu 里编译不要在 Windows 原生 PowerShell 里编译。先设置 RTX3050 的 CUDA 架构减少编译压力export TORCH_CUDA_ARCH_LIST8.6然后安装causal-conv1dcd ~/projects/Vim/causal-conv1d pip install -e . --no-build-isolation再安装mamba-1p1p1cd ~/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 timm0.4.12 einops0.7.0 transformers4.35.2 datasets2.15.0 mlflow2.9.1 opencv-python4.8.1.78 pillow10.1.0 scipy1.11.4 scikit-learn1.3.2 pandas2.1.3 numpy1.26.2 pyyaml6.0.1 tqdm4.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, pretrainedFalse, num_classes1000, ).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 上的通用视觉 backboneHugging Face 模型卡说明它用于分类、检测、分割等视觉任务。创建权重目录cd ~/projects/Vim mkdir -p checkpoints下载 Vim-tinywget -O checkpoints/vim_t_midclstok_76p1acc.pth https://huggingface.co/hustvl/Vim-tiny-midclstok/resolve/main/vim_t_midclstok_76p1acc.pth?downloadtrue如果你想下微调后的 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?downloadtrueVim-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, pretrainedFalse, num_classes1000, ) checkpoint torch.load(ckpt_path, map_locationcpu) state_dict checkpoint[model] if model in checkpoint else checkpoint missing, unexpected model.load_state_dict(state_dict, strictFalse) 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(dim1).item()) if __name__ __main__: main()运行python quick_test_vim_pretrained.py能输出torch.Size([1, 1000])就说明预训练权重也加载成功了。十二、把项目接入 PyCharmPyCharm 官方文档说明可以通过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。