VisionMamba复现-手把手教你安装Causal-Conv1d和Mamba-ssm

最新推荐文章于 2025-03-17 21:00:00 发布

十级钳工郭叔

最新推荐文章于 2025-03-17 21:00:00 发布

阅读量1.2k

点赞数 13

文章标签： python 人工智能深度学习

本文链接：https://blog.csdn.net/weixin_58881174/article/details/144310318

版权

基础环境配置

CUDA==11.8 #好像 Causal-Conv1d 和 Mamba-ssm 支持的最低CUDA版本就是11.8

python==3.9.0

pytorch==2.2.0

torchaudio==2.2.0

torchvision==0.17.0

Causal-Conv1d安装

通常pip install causal-conv1d==1.0.0会遇到各种各样的问题，建议下载源码安装，步骤如下：

git clone https://github.com/Dao-AILab/causal-conv1d.git
cd causal-conv1d
git checkout v1.0.0

git checkout v1.0.0 这一步是切换causal-conv1d的版本，建议去官网看一下causal-conv1d的版本对应的cuda和torch版本以免版本对应不上。

官网连接：

Releases · Dao-AILab/causal-conv1dhttps://github.com/Dao-AILab/causal-conv1d/releases然后执行：

CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install .

最后的 . 要有。

Mamba-ssm安装

与Causal-Conv1d的安装一样，直接下载源码安装：

git clone https://github.com/state-spaces/mamba.git
cd mamba
git checkout v1.0.1

同样这里建议看一下Mamba-ssm对应的cuda和torch版本，官网连接：

Releases · state-spaces/mamba · GitHubhttps://github.com/state-spaces/mamba/releases

这里要修改一下源码

setup.py中修改：

FORCE_BUILD = os.getenv("MAMBA_FORCE_BUILD", "FALSE") == "FALSE"
SKIP_CUDA_BUILD = os.getenv("MAMBA_SKIP_CUDA_BUILD", "FALSE") == "FALSE"

mamba_ssm/ops/selective_scan_interface.py 中修改：

第一步：注释掉一个import语句

# import selective_scan_cuda

第二步：将下面两个函数

def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
                     return_last_state=False):
    """if return_last_state is True, returns (out, last_state)
    last_state has shape (batch, dim, dstate). Note that the gradient of the last state is
    not considered in the backward pass.
    """
    return SelectiveScanFn.apply(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
 
 
def mamba_inner_fn(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    out_proj_weight, out_proj_bias,
    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
    C_proj_bias=None, delta_softplus=True
):
    return MambaInnerFn.apply(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
                              out_proj_weight, out_proj_bias,
                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

改为：

def selective_scan_fn(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False,
                     return_last_state=False):
    """if return_last_state is True, returns (out, last_state)
    last_state has shape (batch, dim, dstate). Note that the gradient of the last state is
    not considered in the backward pass.
    """
    return selective_scan_ref(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
 
def mamba_inner_fn(
    xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
    out_proj_weight, out_proj_bias,
    A, B=None, C=None, D=None, delta_bias=None, B_proj_bias=None,
    C_proj_bias=None, delta_softplus=True
):
    return mamba_inner_ref(xz, conv1d_weight, conv1d_bias, x_proj_weight, delta_proj_weight,
                              out_proj_weight, out_proj_bias,
                              A, B, C, D, delta_bias, B_proj_bias, C_proj_bias, delta_softplus)

然后执行：