search

hustvl / Vim

Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
Apache License 2.0
2.55k stars 159 forks source link

ValueError: Cannot find backend for cpu #56

Open LilLouis5 opened 3 months ago

LilLouis5 commented 3 months ago

when i use a test example in vision_mamba.py，it has this problem. How can i fix it?

besides, the GPT say that this problem is because of the version of Triton , but i check my environment the version is 2.1.0.

Thanks.

Lynxsx commented 2 months ago

i found this problem too, and i put the model to gpu, then it works, maybe this model can't run in cpu

poult-lab commented 2 months ago

i found this problem too, and i put the model to gpu, then it works, maybe this model can't run in cpu

Hello friend, how did you put the model into GPU, can you show us the code?

poult-lab commented 2 months ago

I think I solved it, I just use the below code to put model into GPUs.

`import torch from models_mamba import VisionMamba from timm.models.vision_transformer import VisionTransformer, _cfg

DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu' print("This is the running device:", DEVICE)

input = torch.rand(1, 224, 224).cuda() model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=True, use_middle_cls_token=True) model.default_cfg = _cfg() model.to(DEVICE)

output = model(input) print(output.shape)`