安装VMamba并测试VSSBlock

最近想试一下VMamba的效果，所以想着安装一下VMamba，试一下VSSBlock的效果，这里记录一下安装过程，以备后续查阅。

安装VMamba

step 1：创建一个conda的虚拟环境：

conda create -n vmamba python=3.10
conda activate vmamba

step 2：安装pytorch

conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.8 -c pytorch -c nvidia

step 3：安装causal-conv1d

https://github.com/Dao-AILab/causal-conv1d/releases

pip install causal_conv1d-1.4.0+cu118torch2.0cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

step 4：安装mamba-ssm

pip3 install mamba-ssm

step 5：安装vmamba的依赖项

pip3 install -r requirements.txt

requirements.txt文件的内容

packaging
triton
timm==0.4.12
pytest 
chardet 
yacs 
termcolor
submitit 
tensorboardX
fvcore
seaborn

step 6：安装vmamba

cd kernels/selective_scan && pip3 install .

step 7：解决bug

安装完之后运行出现如下的错误：

NameError: name ‘selective_scan_cuda_core’ is not defined

然后再下面的回答中找到了解决方法，作者说selective_scan_cuda_core已经被弃用了。

NameError: name ‘selective_scan_cuda_core’ is not defined #76

ModuleNotFoundError: No module named ‘selective_scan_cuda’ #55

然后就把下面这句话改掉：

v2=partial(self.forward_corev2, force_fp32=(not self.disable_force32), selective_scan_backend="core"),

v2=partial(self.forward_corev2, force_fp32=(not self.disable_force32), selective_scan_backend="oflex"),

step 8：测试

device = torch.device("cuda:0")
hidden_dim = 96  # 假设 VSSBlock 的隐藏维度为 64
network = VSSBlock(hidden_dim).to(device)
# 随机生成一张输入图片作为示例
input_image = torch.randn(1, 3, 224, 224)  # 假设输入图片大小为 224x224，通道数为 3（RGB图像）
input_image = input_image.to(device)
# 前向传播
patch_embed = make_patch_embed(in_chans=3, embed_dim=96, patch_size=4, patch_norm=True, norm_layer=nn.LayerNorm, channel_first=False).to(device)
pos_embed = make_pos_embed(embed_dims=96, patch_size=4, img_size=224).permute(0, 2, 3, 1).to(device)
input_image = patch_embed(input_image)
input_image = input_image + pos_embed
output = network(input_image)
# 输出结果的大小
print("Output shape:", output.shape)

输出结果：

Output shape: torch.Size([1, 56, 56, 96])

这样就成功了！！！

其他Bug

在安装调试过程中还出现了如下的错误，解决方法如下：

ImportError: libc10.so: cannot open shared object file: No such file or directory

参考资料

https://github.com/state-spaces/mamba

Linux下VMamba 环境复现+环境测试

复现mamba时配置环境的踩坑记录汇总（已成功复现）

https://github.com/state-spaces/mamba/releases

https://github.com/MzeroMiko/VMamba

Failed to build causal-conv1d – 离线安装(mamba_ssm)

‘import selective scan’ 操作导入失败，无法导入 #2