# Write a GPU/CPU agnostic package with Flux

**URL:** <https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [October 10, 2019, 3:26pm UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734 "2019-10-10T15:26:44Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![jonathan-laurent](https://avatars.discourse-cdn.com/v4/letter/j/ecae2f/32.png) [@jonathan-laurent](https://discourse.julialang.org/u/jonathan-laurent)\
**Post date:** [October 10, 2019, 3:26pm UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734/1 "2019-10-10T15:26:44Z")

</div>

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](https://github.com/FluxML/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:

```julia
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:

```julia
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:

```julia-repl
[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!**

---

<div class="post-metadata">

**Author:** ![rapus95](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rapus95/32/3773_2.png) [@rapus95](https://discourse.julialang.org/u/rapus95)\
**Post date:** [October 10, 2019, 3:58pm UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734/2 "2019-10-10T15:58:34Z")

</div>

Right now if you understand you correctly you want some pointers on how to allow switching of resources after writing algorithms (the code for the algorithms doesn’t care about gpu or not anyway according to that: [link](https://fluxml.ai/Flux.jl/v0.4/gpu.html)). Then two packages come to my mind.

1. [Plots.jl](https://github.com/JuliaPlots/Plots.jl) which is somewhat like the general interface to many plotting backends.
2. [ComputationalResources.jl](https://github.com/timholy/ComputationalResources.jl) which is used in Images IIRC to decide whether computations shall be threaded, single core, multicore or GPU.

---

<div class="post-metadata">

**Author:** ![DrChainsaw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/drchainsaw/32/8497_2.png) [@DrChainsaw](https://discourse.julialang.org/u/DrChainsaw)\
**Post date:** [October 10, 2019, 4:04pm UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734/3 "2019-10-10T16:04:36Z")

</div>

Flux exports the function ‘gpu’ which wraps all arrays in CuArrays if cuda is available.

It has a (not exported iirc) function called has\_cuda (or maybe has\_gpu) which it uses to determine if cuda is available. It basically uses the method from this thread to do so [Weak dependence on CUDA packages, 2019 version](https://discourse.julialang.org/t/weak-dependence-on-cuda-packages-2019-version/29033)

---

<div class="post-metadata">

**Author:** ![jonathan-laurent](https://avatars.discourse-cdn.com/v4/letter/j/ecae2f/32.png) [@jonathan-laurent](https://discourse.julialang.org/u/jonathan-laurent)\
**Post date:** [October 10, 2019, 4:20pm UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734/4 "2019-10-10T16:20:38Z")

</div>

The thread you linked to addresses my issue perfectly. Thanks!

---

<div class="post-metadata">

**Author:** ![xiaodai](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/xiaodai/32/15937_2.png) [@xiaodai](https://discourse.julialang.org/u/xiaodai)\
**Post date:** [October 11, 2019, 11:07am UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734/5 "2019-10-11T11:07:02Z")

</div>

> [@jonathan-laurent](#):
>
> namely an implementation of Deepmind’s AlphaZero algorithm

Exciting! Repo? I want to contribute

---

<div class="post-metadata">

**Author:** ![jonathan-laurent](https://avatars.discourse-cdn.com/v4/letter/j/ecae2f/32.png) [@jonathan-laurent](https://discourse.julialang.org/u/jonathan-laurent)\
**Post date:** [October 11, 2019, 1:34pm UTC](https://discourse.julialang.org/t/write-a-gpu-cpu-agnostic-package-with-flux/29734/6 "2019-10-11T13:34:56Z")

</div>

Thanks for your interest!  
The package is still under heavy development but you can already have a look at the [repo](https://github.com/jonathan-laurent/AlphaZero.jl). Hopefully, I’ll be able to make a public release with working examples and proper documentation soon. 🙂