I am trying to setup my machine learning training loop and I don't want it to use the GPU on my computer. I have CUDA setup and installed but I want to run my code on the CPU instead. How do I do this in Flux.jl?
1 Answers
By default, Flux will not use the GPU unless you tell it to. First, ensure CUDA is not loaded my commenting out using CUDA. Next, remove any calls to gpu. In general, if you expect to be selectively running a program either on the CPU and GPU, then write your code using the following design pattern:
using CUDA
CUDA.allowscalar(false)
using Flux
# configure device
device = gpu
# made up function that loads your dataset
data = loaddata() .|> device
# made up function that builds your model
m = mymodel() |> device
# rest of your flux code
Notice how I used the generic device function everywhere so that I only need to set my chosen target device once atstart of the program. If your program is inside a “main” function, you could have that function accept a device argument.
If you absolutely need to disable the GPU, then CUDA.jl looks for the CUDA_VISIBLE_DEVICES environment variable which you can set to be empty. But you would only need to do this for debug purposes/sanity checks.
Related