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