MPS Support

[trbutler]

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]

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

[ProGamerGov]

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

[trbutler]

@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]

@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]

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]

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...

neural-style-pt/CaffeLoader.py

Lines 113 to 120 in 639fa26

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

...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