[教程] Linux系统Anaconda平台下基于conda配置Mamba_ssm环境

写在前面:

• 系统配置：

        Ubuntu20.04 LTS, AMD EPYC 7352, NVIDIA-A100

• 注意事项:

        使用该教程可安装绝大多数版本的 mamba_ssm 和 causal_conv1d, 但是不同版本库的使用会存在细微区别，动手能力强的选手可以自己调整，不能调整的就按照原始代码的 requirements 安装对应版本即可。

省流版本：

1. 第一步就联合 torch、torchvision 一起安装 timm 库; 
2. mamba_ssm 和 causal-conv1d 库直接 pip install 大概率报错;
3. 使用 git clone 和 pip install -e 安装也有很大机率报错;
4. 手动下载 whl (轮子) 再用 pip 安装成功率最高，也最省事
5. 能用 conda 安装就用 conda, 实在没辙再用 pip.

安装流程：

第1步: 初始化环境，联合安装 timm + pytorch 库:

• 注意不要分多次安装，一定要一次性安装 (!!非常重要!!)：

  conda create -n mamba
  conda activate mamba
  
  # timm 包含 pytorch, 安装时间相对较长
  conda install -c pytorch -c conda-forge -c nvidia timm==0.6.5 pytorch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 pytorch-cuda=11.8

第2步: 配置 mamba_ssm 和 causal-conv1d 库：

• 首先查看当前环境的 python, cudatoolkit, pytorch 的版本:

  # 查看 python 版本
  python -V
  # Python    3.9.19
  
  # 查看 cuda 版本
  pip list | grep cuda
  # nvidia-cuda-runtime-cu12    12.4.127
  
  # 查看 torch 版本
  pip list | grep torch
  # torch    2.5.1

• 直接去 mamba_ssm 和 causal-conv1d 的仓库找相应的 whl 包并下载，例如这里需要下载： 

mamba_ssm-2.2.4+cu12torch2.5cxx11abiTRUE-cp39-cp39-linux_x86_64.whlhttps://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.5cxx11abiTRUE-cp39-cp39-linux_x86_64.whl
causal_conv1d-1.5.0.post8+cu12torch2.5cxx11abiTRUE-cp39-cp39-linux_x86_64.whlhttps://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/causal_conv1d-1.5.0.post8+cu12torch2.5cxx11abiTRUE-cp39-cp39-linux_x86_64.whl

• 将上述 whl 包拷贝到相应路径下，安装(安装完毕可删除)：

# /dir/ 表示目标文件路径
pip install /dir/​mamba_ssm-2.2.4+cu11torch2.3cxx11abiTRUE-cp312-cp312-linux_x86_64.whl
pip install /dir/causal_conv1d-1.5.0.post8+cu11torch2.3cxx11abiTRUE-cp312-cp312-linux_x86_64.whl

第3步: 补充安装 triton 库（重要! mamba_ssm 的依赖包，不知为何没自动装上）:

conda install triton

阶段完成: 至此，基础的 torch, mamba_ssm 和 timm 库已经配置完成:

可以进入 python 环境导入相关包，不报错即为成功：

import timm
import mamba_ssm
import causal_conv1d

进阶安装:

此处主要进一步配置 mamba_ssm 环境，补充安装一些项目过程中的必要依赖库 (opencv, open3d及其conda频道, pytorch3d 等等)。读者们可根据实际要求 following 安装或 customizing 安装。

其中，部分库的安装既可以走 conda 频道，也可以走 pypi 频道，这里给出了推荐的优先级，rank后的数字越小优先级越高 (选择一种方式安装即可，切勿重复安装!!).

# 安装相应科学计算库和文件操作库
conda install pyyaml h5py matplotlib scikit-learn scikit-image easydict

# 安装日志和进度条库 (可选)
conda install tqdm wandb

# 安装 pytorch3d 库 (可选)
conda install -c conda-forge -c fvcore -c iopath fvcore iopath
conda install -c pytorch3d pytorch3d # pytorch3d 频道 (rank #1)
conda install -c conda-forge pytorch3d # conda-forge 频道 (rank #2)

# 安装 open3d 库 (可选)
conda install -c conda-forge open3d # conda-forge 频道 (rank #1)
conda install -c open3d-admin open3d # open3d-admin 频道 (rank #2)
pip install open3d # pypi 频道 (rank #3)

# 安装 opencv 库
conda install -c conda-forge opencv # conda 频道 (rank #1)
pip install opencv-python # pypi 频道 (rank #2)

# 备份 mamba_ssm 库文件, 将 path 替换为自己的路径 (可选)
conda env export > /path/environment_mamba_241211.yaml
pip freeze > /path/requirements_mamba_241211.txt

总结:

可以看到，除了mamba_ssm 和 causal-conv1d 必须用基于 whl 文件的离线 pip 安装方法外，本文实现了相对纯净的 conda 式安装 torch + mamba_ssm + pytorch3d + open3d +opencv 以及相关依赖库的方法。

欢迎交流，敬请指正。

转载请联系。