ZHKKKe / MODNet

A Trimap-Free Portrait Matting Solution in Real Time [AAAI 2022]
Apache License 2.0
3.84k stars 636 forks source link

Run locally on CPU? #86

Closed catscratchedme closed 3 years ago

catscratchedme commented 3 years ago

I would like to run the image matting demo locally on my laptop, but I do not have a dedicated GPU.

When running Step 3: Inference, I get the following errors: File "/home/*/.local/lib/python3.8/site-packages/torch/cuda/init.py", line 170, in _lazy_init torch._C._cuda_init() RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

Can this script be modified to run on the CPU only? For those without a dedicated GPU or with an AMD GPU.

ZHKKKe commented 3 years ago

Hi, thanks for your attention.

Please try to delete all .cuda() in the code.

kaisark commented 3 years ago

I was able to run the Pytorch code using the pretrained model and modifying (removing cuda references, add cpu) in the MODNet/torchscript/export_torchscript.py file:

modnet = nn.DataParallel(modnet)
state_dict = torch.load(args.ckpt_path,map_location=torch.device('cpu'))