解决causal_conv1d和mamba_ssm无法安装 -> 直接使用Mamba基础环境docker镜像

介绍

Mamba: Linear-Time Sequence Modeling with Selective State Spaces
论文:https://arxiv.org/abs/2312.00752
Code:https://github.com/state-spaces/mamba 目前 5.3k Star
同时适配VMamba

主要是为了解决pip install causal_conv1d mamba_ssm 老是安装出错。
包括但不限于各种 setup 不了,版本对不上…:

# 诸如
Building wheel for causal-conv1d (setup.py) ... error
error: subprocess-exited-with-error

所以直接建了个docker,但是还是有问题,最后按照 Mamba repo issue#55 的方法总算是装上了。
包含Mamba v1.1.1 的基础环境,Ubuntu20.04,是以nvidia/cuda官方11.7版本docker为基础构建的,具体有以下包:

cuda==11.7 (docker: nvidia/cuda)

pytorch==1.13

causal_conv1d==1.1.1

mamba_ssm==1.1.1

使用

DockHub仓库地址:https://hub.docker.com/repository/docker/kom4cr0/cuda11.7-pytorch1.13-mamba1.1.1/general

docker pull kom4cr0/cuda11.7-pytorch1.13-mamba1.1.1:1.1.1 
# 这个docker的tag是1.1.1,请不要使用latest或者默认

测试

python  # 或者python3
>>>import torch
>>>from mamba_ssm import Mamba
  • 17
    点赞
  • 42
    收藏
    觉得还不错? 一键收藏
  • 32
    评论
评论 32
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值