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!