Windows系统中Python环境下的Mamba_SSM安装流程简介

mamba_ssm官方源码只提供了linux上whl文件并没有提供windows上提供的whl,官方并不支持windows安装,需要修改源码自己安装才能实现。因此在windows上安装有点麻烦。因此提供一个不通过源码编译方式安装,并最终测试安装效果。

因此需要提供下面whl文件进行安装


numpy-1.24.1-cp310-cp310-win_amd64.whl
torch-2.1.1+cu118-cp310-cp310-win_amd64.whl
torchaudio-2.1.1+cu118-cp310-cp310-win_amd64.whl
torchvision-0.16.1+cu118-cp310-cp310-win_amd64.whl
causal_conv1d-1.1.1-cp310-cp310-win_amd64.whl
mamba_ssm-1.1.3-cp310-cp310-win_amd64.whl 

triton-2.1.0-cp310-cp310-win_amd64.whl

其中有些模块只需要pip安装即可,最终只需要准备3个whl文件即可,文件可以网上找一下,也可以去gitee.com/FIRC/pythonlibs_whl_mirror找一下

triton-2.1.0-cp310-cp310-win_amd64.whl

causal_conv1d-1.1.1-cp310-cp310-win_amd64.whl
mamba_ssm-1.1.3-cp310-cp310-win_amd64.whl 

找到上面文件后我们开始:

conda create -n py310 python=3.10 -y

conda activate py310

pip install numpy==1.24.1 setuptools==68.2.2
pip install triton-2.1.0-cp310-cp310-win_amd64.whl
pip install torch==2.1.1+cu118 torchvision==0.16.1+cu118 torchaudio==2.1.1+cu118 -f https://download.pytorch.org/whl/torch_stable.html
pip install causal_conv1d-1.1.1-cp310-cp310-win_amd64.whl
pip install  mamba_ssm-1.1.3-cp310-cp310-win_amd64.whl

上面步骤完成后即可安装完成。最后使用下面代码测试

import torch
from mamba_ssm import Mamba
 
batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape
print('success')

作者:FL1623863129

物联沃分享整理
物联沃-IOTWORD物联网 » Windows系统中Python环境下的Mamba_SSM安装流程简介

发表回复