Windows系统本地部署Vision_Mamba完美适配方案
2026/6/1 6:18:26 网站建设 项目流程

以下方案为个人试错后经验总结,在下面部署流程中可能遇到若干问题可与主页另一篇《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_conv1dmamba-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_resourcessetup.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,但它包含较多检测、分割、日志相关依赖,例如mmcvmmsegmentation等。你现在如果只是先部署并跑通 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.pthvim_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改成21

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询