example for using apple GPU with flux

[rcherukuri12]

Motivation and description

Need an example for using Apple silicon.

Need clear documentation that tell you to do things twice.

First time :

Using Flux, Metal
Flux.gpu_backend!( "Metal")

Then exit session

Then launch again

Using Flux, Metal
Flux.GPU_BACKEND
Flux.GPU_BACKEND = 'Metal'

Possible Implementation

Document procedure properly in the front page for GPU switching .

[mcabbott]

There's some documentation here: http://fluxml.ai/Flux.jl/stable/gpu/#Selecting-GPU-backend Could always be better though.

Support for anything but CUDA is a bit experimental. For me e.g. this model http://fluxml.ai/Flux.jl/stable/models/quickstart/ does not work on Metal, but I didn't try hard & maybe have wrong versions or something. Do some models work for you?

I believe that many functions in NNlib need to call the Metal equivalents, parallel to how things here https://github.com/FluxML/NNlib.jl/tree/master/ext/NNlibCUDAExt call CUDA functions.

[rcherukuri12]

Thank you .
For example the quick start example works perfectly with Apple silicon but needs to be documented.

The only thing I had to do is switch to "using Metal" instead of "CUDA"
But also exit and create new session . This allowed re-compilation of Flux with Metal.

Can I recommend you put comments in the code ?
From this :

using Flux, CUDA, Statistics, ProgressMeter
Like this :
using CUDA # switch to 'using Metal' for Apple silicon
using Flux, Statistics, ProgressMeter

[christiangnrd]

Can I recommend you put comments in the code ?

Feel free to submit a pull request! However, it takes 95 seconds to train on my Apple Silicon GPU and 7 seconds to train on my Pi5's CPU so it might be worth waiting for better support before recommending it in the documentation.

[rcherukuri12]

Thank you, will do a pull request.
I would think the other way , the more people start using, the more requests you get ( know the demand for using Flux on Mac with GPU ), so more incentive to make it faster, quicker.

[CarloLucibello]

as with any other backend, and as documented, in order to use metal one has to simply create a metal device and move model and data with that:

using Flux, Metal
device = gpu_device()
model = ... |> device