helper function for selecting a gpu in multi-gpu setting

[CarloLucibello]

As per discussion in https://github.com/JuliaGPU/CuArrays.jl/issues/52, that I report here, the current way to use a specific gpu in a Flux script would be the following

gpu_id = 0  ## set < 0 for no cuda, >= 0 for using a specific device (if available)

if CUDAapi.has_cuda_gpu() && gpu_id >=0
    CUDAdrv.device!(gpu_id)
    CuArrays.allowscalar(false)
    device = Flux.gpu
    @info "Training on GPU-$(gpu_id)"
else
    device = Flux.cpu
    @info "Training on CPU"
end

model = model |> device
for x in data
    x = x |> device
    ....

This involves importing 3 cuda libraries, CUDAapi, CuArrays and CUDAdrv, which likely won't be likely used in the rest of the script.

I suggest to wrap the same functionality under Flux, in a gpu! function, so that we can equivalently write

gpu_id = 0  ## set < 0 for no cuda, >= 0 for using a specific device (if available)

if CUDAapi.has_cuda_gpu() && gpu_id >=0
    device = Flux.gpu!(gpu_id, allowscalar=false)
    @info "Training on GPU-$(gpu_id)"
else
    device = Flux.cpu
    @info "Training on CPU"
end

Thoughts?

[DhairyaLGandhi]

CuArrays.device!(id) do
  ... # the usual
end

Should be equivalent and cleaner

See also: FluxML/Hydra.jl

[CarloLucibello]

you mean wrapping model creation and training within the do block?

[DhairyaLGandhi]

Yeah, most times you want to be data parallel, so having multiple workers start tasks on different gpus and syncing appropriately is usually preferred.

[CarloLucibello]

See also: FluxML/Hydra.jl

what should I look at?

[DhairyaLGandhi]

It's way to move model training keepings things in sync has been useful in the past to avoid waiting around too long when you have heterogenous devices in the same machine

[CarloLucibello]

multi-gpu training and data parallelism is something we definitely want, but this issue is about the much simpler problem of selecting a gpu for single gpu training in a multi-gpu setting.
For that, it just seems convenient to have Flux.gpu!(gpu_id).

[johnnychen94]

Flux.device!(CUDALibs(gpu_id))
Flux.device!(CPU1())

looks more intuitive to me

https://github.com/timholy/ComputationalResources.jl/blob/master/src/ComputationalResources.jl#L127-L129

[CarloLucibello]

Utilities in MLDataDevices.jl can take care of this.