Windows系统本地部署Vision_Mamba完美适配方案

m0_63615578

370人浏览 · 2026-05-31 01:41:48

m0_63615578 · 2026-05-31 01:41:48 发布

以下方案为个人试错后经验总结，在下面部署流程中可能遇到若干问题可与主页另一篇《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。