You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
@@ -16,68 +16,13 @@ in your code. Notice that for CUDA, explicitly loading also `cuDNN` is not requi
16
16
!!! compat "Flux ≤ 0.13"
17
17
Old versions of Flux automatically installed CUDA.jl to provide GPU support. Starting from Flux v0.14, CUDA.jl is not a dependency anymore and has to be installed manually.
18
18
19
-
## Checking GPU Availability
20
-
21
-
By default, Flux will run the checks on your system to see if it can support GPU functionality. You can check if Flux identified a valid GPU setup by typing the following:
22
-
23
-
```julia
24
-
julia>using CUDA
25
-
26
-
julia> CUDA.functional()
27
-
true
28
-
```
29
-
30
-
For AMD GPU:
31
-
32
-
```julia
33
-
julia>using AMDGPU
34
-
35
-
julia> AMDGPU.functional()
36
-
true
37
-
38
-
julia> AMDGPU.functional(:MIOpen)
39
-
true
40
-
```
41
-
42
-
For Metal GPU:
43
-
44
-
```julia
45
-
julia>using Metal
46
-
47
-
julia> Metal.functional()
48
-
true
49
-
```
50
-
51
-
## Selecting GPU backend
52
-
53
-
Available GPU backends are: `CUDA`, `AMDGPU` and `Metal`.
54
-
55
-
Flux relies on [Preferences.jl](https://github.com/JuliaPackaging/Preferences.jl) for selecting default GPU backend to use.
56
-
57
-
There are two ways you can specify it:
58
-
59
-
- From the REPL/code in your project, call `Flux.gpu_backend!("AMDGPU")` and restart (if needed) Julia session for the changes to take effect.
60
-
- In `LocalPreferences.toml` file in you project directory specify:
61
-
```toml
62
-
[Flux]
63
-
gpu_backend = "AMDGPU"
64
-
```
65
-
66
-
Current GPU backend can be fetched from `Flux.GPU_BACKEND` variable:
67
-
68
-
```julia
69
-
julia> Flux.GPU_BACKEND
70
-
"CUDA"
71
-
```
72
-
73
-
The current backend will affect the behaviour of methods like the method `gpu` described below.
74
19
75
20
## Basic GPU Usage
76
21
77
22
Support for array operations on other hardware backends, like GPUs, is provided by external packages like [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl), [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl), and [Metal.jl](https://github.com/JuliaGPU/Metal.jl).
78
23
Flux is agnostic to array types, so we simply need to move model weights and data to the GPU and Flux will handle it.
79
24
80
-
For example, we can use `CUDA.CuArray` (with the `cu` converter) to run our [basic example](@ref man-basics) on an NVIDIA GPU.
25
+
For example, we can use `CUDA.CuArray` (with the `CUDA.cu` converter) to run our [basic example](@ref man-basics) on an NVIDIA GPU.
81
26
82
27
(Note that you need to have CUDA available to use CUDA.CuArray – please see the [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) instructions for more details.)
83
28
@@ -146,6 +91,50 @@ julia> x |> cpu
146
91
0.7766742
147
92
```
148
93
94
+
## Using device objects
95
+
96
+
In Flux, you can create `device` objects which can be used to easily transfer models and data to GPUs (and defaulting to using the CPU if no GPU backend is available).
97
+
These features are provided by [MLDataDevices.jl](https://github.com/LuxDL/MLDataDevices.jl) package, that Flux uses internally and re-exports.
98
+
99
+
Device objects can be automatically created using the [`cpu_device`](@ref MLDataDevices.cpu_device) and [`gpu_device`](@ref MLDataDevices.gpu_device) functions. For instance, the `gpu` and `cpu` functions are just convenience functions defined as
100
+
101
+
```julia
102
+
cpu(x) =cpu_device()(x)
103
+
gpu(x) =gpu_device()(x)
104
+
```
105
+
106
+
`gpu_device` performs automatic GPU device selection and returns a device object:
107
+
- If no GPU is available, it returns a `CPUDevice` object.
108
+
- If a LocalPreferences file is present, then the backend specified in the file is used. To set a backend, use `Flux.gpu_backend!(<backend_name>)`. If the trigger package corresponding to the device is not loaded (e.g. with `using CUDA`), then a warning is displayed.
109
+
- If no LocalPreferences option is present, then the first working GPU with loaded trigger package is used.
110
+
111
+
Consider the following example, where we load the [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) package to use an NVIDIA GPU (`"CUDA"` is the default preference):
112
+
113
+
```julia-repl
114
+
julia> using Flux, CUDA;
115
+
116
+
julia> device = gpu_device() # returns handle to an NVIDIA GPU if available
117
+
(::CUDADevice{Nothing}) (generic function with 4 methods)
118
+
119
+
julia> model = Dense(2 => 3);
120
+
121
+
julia> model.weight # the model initially lives in CPU memory
122
+
3×2 Matrix{Float32}:
123
+
-0.984794 -0.904345
124
+
0.720379 -0.486398
125
+
0.851011 -0.586942
126
+
127
+
julia> model = model |> device # transfer model to the GPU
128
+
Dense(2 => 3) # 9 parameters
129
+
130
+
julia> model.weight
131
+
3×2 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
132
+
-0.984794 -0.904345
133
+
0.720379 -0.486398
134
+
0.851011 -0.586942
135
+
```
136
+
137
+
149
138
## Transferring Training Data
150
139
151
140
In order to train the model using the GPU both model and the training data have to be transferred to GPU memory. Moving the data can be done in two different ways:
@@ -227,65 +216,8 @@ To select specific devices by device id:
227
216
$ export CUDA_VISIBLE_DEVICES='0,1'
228
217
```
229
218
230
-
231
219
More information for conditional use of GPUs in CUDA.jl can be found in its [documentation](https://cuda.juliagpu.org/stable/installation/conditional/#Conditional-use), and information about the specific use of the variable is described in the [Nvidia CUDA blog post](https://developer.nvidia.com/blog/cuda-pro-tip-control-gpu-visibility-cuda_visible_devices/).
232
220
233
-
## Using device objects
234
-
235
-
As a more convenient syntax, Flux allows the usage of GPU `device` objects which can be used to easily transfer models to GPUs (and defaulting to using the CPU if no GPU backend is available). This syntax has a few advantages including automatic selection of the GPU backend and type stability of data movement.
236
-
These features are provided by [MLDataDevices.jl](https://github.com/LuxDL/MLDataDevices.jl) package, that Flux's uses internally and re-exports.
237
-
238
-
A `device` object can be created using the [`gpu_device`](@ref MLDataDevices.gpu_device) function.
239
-
`gpu_device` first checks for a GPU preference, and if possible returns a device for the preference backend. For instance, consider the following example, where we load the [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) package to use an NVIDIA GPU (`"CUDA"` is the default preference):
240
-
241
-
```julia-repl
242
-
julia> using Flux, CUDA;
243
-
244
-
julia> device = gpu_device() # returns handle to an NVIDIA GPU if available
245
-
(::CUDADevice{Nothing}) (generic function with 4 methods)
246
-
247
-
julia> model = Dense(2 => 3);
248
-
249
-
julia> model.weight # the model initially lives in CPU memory
250
-
3×2 Matrix{Float32}:
251
-
-0.984794 -0.904345
252
-
0.720379 -0.486398
253
-
0.851011 -0.586942
254
-
255
-
julia> model = model |> device # transfer model to the GPU
256
-
Dense(2 => 3) # 9 parameters
257
-
258
-
julia> model.weight
259
-
3×2 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
260
-
-0.984794 -0.904345
261
-
0.720379 -0.486398
262
-
0.851011 -0.586942
263
-
```
264
-
265
-
The device preference can also be set via the [`gpu_backend!`](@ref MLDataDevices.gpu_backend!) function. For instance, below we first set our device preference to `"AMDGPU"`:
266
-
267
-
```julia-repl
268
-
julia> gpu_backend!("AMDGPU")
269
-
[ Info: GPU backend has been set to AMDGPU. Restart Julia to use the new backend.
270
-
```
271
-
If no functional GPU backend is available, the device will default to a CPU device.
272
-
You can also explictly request a CPU device by calling the [`cpu_device`](@ref MLDataDevices.cpu_device) function.
273
-
274
-
```julia-repl
275
-
julia> using Flux, MLDataDevices
276
-
277
-
julia> cdev = cpu_device()
278
-
(::CPUDevice{Nothing}) (generic function with 4 methods)
279
-
280
-
julia> gdev = gpu_device(force=true) # force GPU device, error if no GPU is available
281
-
(::CUDADevice{Nothing}) (generic function with 4 methods)
282
-
283
-
julia> model = Dense(2 => 3); # model in CPU memory
284
-
285
-
julia> gmodel = model |> gdev; # transfer model to GPU
286
-
287
-
julia> cmodel = gmodel |> cdev; # transfer model back to CPU
We don't run CUDA-aware tests so you're running it at own risk.
481
395
396
+
397
+
## Checking GPU Availability
398
+
399
+
By default, Flux will run the checks on your system to see if it can support GPU functionality. You can check if Flux identified a valid GPU setup by typing the following:
Flux relies on the [MLDataDevices.jl](https://github.com/LuxDL/MLDataDevices.jl/blob/main/src/public.jl) package to manage devices and transfer data across them. You don't have to explicitly use the package, as Flux re-exports the necessary functions and types.
0 commit comments