ProGamerGov / neural-style-pt

PyTorch implementation of neural style transfer algorithm
MIT License
833 stars 178 forks source link

MPS Support #96

Open trbutler opened 1 year ago

trbutler commented 1 year ago

Running neural-style-pt on Apple Silicon seems to require using only CPU or it terminates with this error:

raise AssertionError("Torch not compiled with CUDA enabled")

AssertionError: Torch not compiled with CUDA enabled

Given that PyTorch now supports native Apple Metal acceleration, is there a way to fix this so it'd use MPS?

ProGamerGov commented 1 year ago

@ProGamerGov Yeah it should be relatively straight forward, however I currently don't have a way to test it.

ProGamerGov commented 1 year ago

@trbutler Try running this WIP neural-style-pt branch on MPS: https://github.com/ProGamerGov/neural-style-pt/tree/master-2

trbutler commented 1 year ago

@ProGamerGov Thanks! I just tried it and it doesn't seem to be paying attention to my request to switch backends. For example upon typing:

python neural_style.py -style_image /Users/timothybutler/Downloads/f24.jpg -content_image /Users/timothybutler/Downloads/SCAN22122008_00000.tif -backend mps

It starts and then gives the same error:

Traceback (most recent call last): File "/Users/timothybutler/Experiments/neural-style-pt/neural_style.py", line 500, in main() File "/Users/timothybutler/Experiments/neural-style-pt/neural_style.py", line 62, in main content_image = preprocess(params.content_image, params.image_size).to(backward_device) File "/Users/timothybutler/Experiments/miniconda3/envs/ldm/lib/python3.10/site-packages/torch/cuda/init.py", line 211, in _lazy_init raise AssertionError("Torch not compiled with CUDA enabled") AssertionError: Torch not compiled with CUDA enabled

Perhaps the backend isn't getting passed along?

trbutler commented 1 year ago

@ProGamerGov I just realized if I set both -backend and -gpu to mps it does seem to work. Should setting backend to mps set gpu to the same in setup_gpu if nothing is specified in the gpu parameter?

Currently running neural_style on an image to see what happens, but it does show 97% GPU usage from the python process in MacOS's Activity Monitor.

trbutler commented 1 year ago

Results using mps vs. cpu:

python neural_style.py -style_image /Users/timothybutler/Downloads/f24.jpg 77.40s user 81.96s system 45% cpu 5:52.32 total python neural_style.py -style_image /Users/timothybutler/Downloads/f24.jpg 856.05s user 451.82s system 261% cpu 8:20.04 total

I couldn't get -gpu mps,cpu to work -- I thought I'd try that out of curiosity, too, but that reports that AssertionError: The number of -multidevice_strategy layer indices minus 1, must be equal to the number of -gpu devices.

e13h commented 1 year ago

I couldn't get -gpu mps,cpu to work -- I thought I'd try that out of curiosity, too, but that reports that AssertionError: The number of -multidevice_strategy layer indices minus 1, must be equal to the number of -gpu devices.

I think this this because the default for -multidevice_strategy is 4,7,29, which implies you are using 4 devices. If using -gpu mps,cpu, try using a setting -multidevice_strategy to a single number (like 7 for example).


However, this leads to another error

RuntimeError: Invalid device string: 'cuda:mps'

But this can be easily resolved by modifying these lines... https://github.com/ProGamerGov/neural-style-pt/blob/639fa264f74cdb061dd4beee480b7f9126ab3f05/CaffeLoader.py#L113-L120 ...to the following:

    def name_devices(self, input_list):
        device_list = []
        for i, device in enumerate(input_list):
            if str(device).lower() not in ("cpu", "mps"):
                device_list.append("cuda:" + str(device))
            else:
                device_list.append(device)
        return device_list