,Ubuntu上一次性跑通Mamba-YOLO!)
1. Ubuntu 22.04 跑 Mamba-YOLO 到底卡在哪Mamba-YOLO 是把 Mamba 状态空间模型塞进 YOLO 检测头里的一个开源工作适合做目标检测又想吃长序列建模红利的同学。它跟普通 YOLO 最大的区别在于多了一个selective_scan的 CUDA 扩展这个扩展不是纯 Python 包编译时对 CUDA 版本、PyTorch 版本、Python 版本三者极其敏感版本错一位就是一堆undefined symbol或者nvcc fatal。我在 Ubuntu 22.04 上从零跑通这套东西前后重装了三次环境最后锁定在 Python 3.12 PyTorch 2.3.0cu121 selective_scan 0.0.2 这个组合上一次装完直接能前向推理。这篇就把完整流程、可复制的 requirements、环境变量、安装命令和报错排查清单都摊开写你照着敲基本不用再走弯路。适合谁看手里有 Ubuntu 22.04物理机或 WSL2 都行、装了 NVIDIA 显卡驱动、想跑通 Mamba-YOLO 单图推理的人。如果你连 conda 都没装先补一下 Miniconda再回来跟着走。核心检索词先摆出来Ubuntu 上跑 Mamba-YOLO关键就是 selective_scan 安装包要跟 Python 3.12 对齐CUDA 要跟 PyTorch 的 cu121 对齐缺一个都跑不起来。2. 前置核对CUDA、驱动、PyTorch 三者对齐2.1 先看驱动和 CUDA 版本打开终端先确认显卡驱动在且 CUDA 版本不低于 12.1nvidia-smi输出右上角会写CUDA Version: 12.x。注意这个数字是驱动支持的最高 CUDA 版本不是你已经装的 CUDA Toolkit 版本。Mamba-YOLO 的 selective_scan 编译需要 nvcc所以还得确认 Toolkitnvcc --version如果提示 command not found说明只装了驱动没装 Toolkit。Ubuntu 22.04 上装 CUDA Toolkit 12.1 可以用官方 runfile也可以用 conda 装cuda-toolkit12.1。我实测下来用 conda 装最省事不会污染系统conda install -c nvidia/label/cuda-12.1.0 cuda-toolkit -y装完再nvcc --version应该显示release 12.1。2.2 创建 Python 3.12 环境selective_scan 的预编译 whl 我只找到 cp312 这一个版本所以 Python 必须是 3.12别用 3.10 或 3.11否则要么找不到 whl要么自己编译踩坑。conda create -n mambayolo python3.12 -y conda activate mambayolo2.3 装 cu121 版 PyTorch这一步版本必须锁死torch 2.3.0、torchvision 0.18.0、torchaudio 2.3.0全部走 cu121 源pip3 install torch2.3.0cu121 torchvision0.18.0cu121 torchaudio2.3.0 \ --index-url https://download.pytorch.org/whl/cu121装完验证一下python -c import torch; print(torch.__version__, torch.cuda.is_available())输出应该是2.3.0cu121 True。如果cuda.is_available()是 False先别往下走回去查驱动和 Toolkit。3. 可复制配置requirements 与环境变量3.1 requirements.txt把下面内容存成requirements.txt放在项目根目录torch2.3.0cu121 torchvision0.18.0cu121 torchaudio2.3.0 seaborn thop timm einops pyyaml tqdm matplotlib opencv-python注意 torch 那三行如果你已经装好了可以注释掉避免 pip 重新解析版本。其余依赖是 Mamba-YOLO 训练和推理都要用的thop算 FLOPstimm提供 backbone 组件einops是 Mamba 模块里做张量重排的。3.2 环境变量selective_scan 编译和运行时会找 CUDA 路径建议在~/.bashrc里加export CUDA_HOME/usr/local/cuda-12.1 export PATH$CUDA_HOME/bin:$PATH export LD_LIBRARY_PATH$CUDA_HOME/lib64:$LD_LIBRARY_PATH export TORCH_CUDA_ARCH_LIST8.6TORCH_CUDA_ARCH_LIST按你的显卡填30 系是 8.640 系是 8.920 系是 7.5。填错会导致编译出来的 kernel 跟显卡不匹配运行时报no kernel image is available。改完source ~/.bashrc。提示如果你用 conda 装的 cuda-toolkitCUDA_HOME 应该指向$CONDA_PREFIX可以用echo $CONDA_PREFIX看一下实际路径再填。4. 装 selective_scan 与 Mamba-YOLO4.1 安装 selective_scan 预编译包这是整个流程最关键的一步。预编译 whl 文件名是selective_scan-0.0.2-cp312-cp312-linux_x86_64.whlcp312 对应 Python 3.12linux_x86_64 对应系统架构。直接pip install selective_scan-0.0.2-cp312-cp312-linux_x86_64.whl装完验证python -c import selective_scan_cuda; print(ok)没报错就说明 CUDA 扩展加载成功。如果报ImportError: libcudart.so.12: cannot open shared object file说明LD_LIBRARY_PATH没配好回去检查 3.2 的环境变量。4.2 克隆并安装 Mamba-YOLOgit clone https://github.com/HZAI-ZJNU/Mamba-YOLO.git cd Mamba-YOLO pip install -v -e .-e是 editable 模式方便你改代码。-v打印详细日志出问题好定位。如果 git clone 卡住可以手动下载 zip 解压重命名为 Mamba-YOLO 再进目录。装完检查一下python -c import mamba_yolo; print(installed)4.3 装剩余依赖pip install -r requirements.txt如果 requirements 里 torch 那几行导致重装直接跳过手动装剩下的pip install seaborn thop timm einops pyyaml tqdm matplotlib opencv-python5. 验证请求单图前向推理跑通5.1 准备权重和测试图从官方仓库 release 或 README 里给的链接下载 Mamba-YOLO 权重比如mamba_yolo_t.pt。测试图随便找一张 COCO 风格的 jpg命名test.jpg。5.2 写一个最小推理脚本新建infer.pyimport torch from PIL import Image from mamba_yolo import build_model device cuda if torch.cuda.is_available() else cpu model build_model(mamba_yolo_t, num_classes80).to(device) model.load_state_dict(torch.load(mamba_yolo_t.pt, map_locationdevice)) model.eval() img Image.open(test.jpg).convert(RGB).resize((640, 640)) x torch.from_numpy(np.array(img)).permute(2, 0, 1).float().unsqueeze(0) / 255.0 x x.to(device) with torch.no_grad(): out model(x) print(output shape:, [o.shape for o in out] if isinstance(out, (list, tuple)) else out.shape)实际模型构建函数名以仓库为准有的版本是build_model有的是MambaYOLO类直接实例化。跑之前先python -c import mamba_yolo; print(dir(mamba_yolo))看一眼导出接口。5.3 运行并看结果python infer.py成功的话会打印输出张量的 shape比如[1, 84, 8400]这种检测头输出。第一次跑会慢一点因为 CUDA kernel 要 JIT 编译第二次就快了。注意如果输出 shape 里 batch 维度是 1说明前向通了。这时候可以接 NMS 后处理画框但那是下一步的事先确认模型能跑。6. 本篇常见错排查清单6.1ImportError: selective_scan_cuda找不到九成是 Python 版本不对。pip list | grep selective看装的是不是 cp312 那个 whl。如果装的是源码编译版检查nvcc --version和CUDA_HOME。6.2undefined symbol: _ZN3c10...PyTorch 版本和 selective_scan 编译时的版本不一致。selective_scan 0.0.2 是对着 torch 2.3.0 编的你装 2.4 或 2.2 都可能出这个错。锁死 2.3.0cu121。6.3no kernel image is available for executionTORCH_CUDA_ARCH_LIST跟显卡算力不匹配。30 系填 8.640 系填 8.9。改完环境变量要重新装 selective_scan因为 kernel 是编译期决定的。6.4CUDA out of memory推理时 batch 设 1图片 resize 到 640。训练才需要大显存单图推理 8G 卡够用。如果还爆检查是不是没加torch.no_grad()。6.5ModuleNotFoundError: No module named mamba_yolopip install -e .没在 Mamba-YOLO 根目录执行或者执行时报错了没注意。重新cd Mamba-YOLO pip install -v -e .看日志最后有没有Successfully installed。6.6 git clone 超时换手动下载 zip解压后重命名目录再进目录pip install -e .。WSL 里如果访问 GitHub 慢可以先在 Windows 下载再拖进 WSL 文件系统。7. 跑通之后把模型对话和编码工作流接上环境跑通只是第一步后面你要调 prompt、试不同检测头配置、对比不同 backbone这些如果每次都本地起服务会很烦。我自己的做法是把模型推理和代码调试分开本地 Ubuntu 只管跑 Mamba-YOLO 的前向和训练模型选型、参数对比、代码片段生成这类事放到 TaoToken 的模型对话里做省得来回切终端。如果你要长期做 Mamba-YOLO 的二次开发比如改 selective_scan 的 kernel、加自定义检测头建议用 Coding Plan 把编码工作流固定下来配合 API Keys 把推理请求脚本化。接入文档里有完整的 curl 和 Python 示例照着改 base_url 就行。具体入口模型对话https://taotoken.net/models?utm_sourcetaotoken_aicg_blog_endutm_contentmodel_chatutm_campaignrewriteCoding Planhttps://taotoken.net/coding-plan?utm_sourcetaotoken_aicg_blog_endutm_contentcoding_planutm_campaignrewriteAPI Keyshttps://taotoken.net/api-keys?utm_sourcetaotoken_aicg_blog_endutm_contentapi_keysutm_campaignrewrite接入文档https://taotoken.net/doc?utm_sourcetaotoken_aicg_blog_endutm_contentdocutm_campaignrewriteClaudeCodeAnthropichttps://taotoken.net/claude-code-anthropic?utm_sourcetaotoken_aicg_blog_endutm_contentclaude_codeutm_campaignrewriteAPI 地址统一用https://taotoken.net/api不要加 UTM 后缀那是给网页链接用的。最后说个实测经验selective_scan 装好之后别急着删 whl 文件换环境或者重装系统时它比源码编译省至少半小时。把 whl、requirements.txt、环境变量三样存一个文件夹下次直接复制粘贴Ubuntu 22.04 上跑通 Mamba-YOLO 就是十分钟的事。