Pytorch MPS Mac M1 Support

[adamdavidcole]

Is your feature request related to a problem? Please describe. I'd like to be able to generate images using the Metal Performance Shaders (MPS) pytorch acceleration for M1 macs

Describe the solution you'd like Run the generate.py script with device=mps

Describe alternatives you've considered I've modified the code and am able to run the function, but the resulting image is completely gray. I'm curious if anyone else has tried and succeeded.

The issue seems to be in w_to_img -> G.synthesis(). The output of that function does not match the output when I run with device=cpu. Up until that point everything matched (for example, the output of get_w_from_seed was correct with mps).

The code changes were:

1. Set device to mps if selected

   if torch.cuda.is_available() and device == 'cuda':
       device = torch.device('cuda')
   elif torch.backends.mps.is_available() and device == "mps":
       device = torch.device('mps')
   else:
       device = torch.device('cpu')

2. Ensure float32 conversion, for example

   z = z.astype(np.float32)
   w = G.mapping(torch.from_numpy(z).to(device), None)

[PDillis]

For running with device='cpu', do you use use_cpu(G) before generating the images? I have an example in generate.py, see if this helps!

I would love to solve this problem myself if only I had an M1 chip. So in the meantime, I'll depend on your results (and others' input)!

[adamdavidcole]

Yea, I experimented with the use_cpu() and the force_fp32() options, but no luck.

The version of pytorch that runs on the M1 gpu is super new, so could be a bug on their side of things. I'll keep an eye on it and maybe it'll resolve itself in a future release 🤞

Curious to see if anyone else comes up with a workaround or solution though!

[PDillis]

Hey! Any updates on the behavior in Mac with the latest updates from Pytorch? Closing this issue for the time being, but am very interested in any additional results.