I am currently developing a Julia package based on Flux (namely an implementation of Deepmind’s AlphaZero algorithm) and I am looking for the best way to make it accessible to people with and without a GPU.
Ideally, enabling GPU usage should be as simple as providing a boolean parameter when calling the training algorithm (an error message would be displayed if no GPU is detected). I am having some trouble to make this happen though.
I looked at the solution used by model-zoo, which consists in providing two distinct project files for each model. Then, a model can be run on GPU by including its source file after running the following snippet:
using Pkg; Pkg.activate("cuda"); Pkg.instantiate()
using CuArrays
I do not see how to use this solution within a package though (as a package cannot have several project files).
I could imagine a solution looking something like this:
module AlphaZero
using CuArrays
using Flux
const gpu = gpu_detected() ? Flux.gpu : identity
# ...
end
Here, the package CuArrays would be imported but never actually used on a machine without a GPU.
There are two problems with this though. First, I am not aware of any function such as gpu_detected. There would be Flux.use_cuda but this variable is set based on whether or not CuArrays is imported. Second, importing CuArrays on a machine without a GPU yields an error:
[ Info: Precompiling CuArrays [3a865a2d-5b23-5a0f-bc46-62713ec82fae]
ERROR: LoadError: Could not find CUDA driver library
I would really appreciate your ideas on how to best write CPU/GPU agnostic packages!
