Поддержка GPU

Начиная с версии Flux 0.14 от пользователя не требуется использовать определенный бэкенд GPU и соответствующие зависимости пакетов. Благодаря механизму расширения пакетов, появившемуся в Julia 1.9, Flux загружает относящийся к GPU код, как только становится доступен пакет GPU (например, посредством using CUDA).

Для поддержки GPU NVIDIA необходимо установить в среде пакеты CUDA.jl и cuDNN.jl. Для их установки введите в REPL Julia команду ] add CUDA, cuDNN. Дополнительные сведения см. в файле сведений CUDA.jl.

Поддержка GPU AMD доступна начиная с версии Julia 1.9 в системах с установленными библиотеками ROCm и MIOpen. Дополнительные сведения см. в репозитории AMDGPU.jl.

Ускорение GPU Metal доступно на оборудовании Apple Silicon. Дополнительные сведения см. в репозитории Metal.jl. Поддержка Metal во Flux находится на экспериментальной стадии, и многие функции пока не доступны.

Чтобы включить поддержку GPU во Flux, необходимо вызвать using CUDA, using AMDGPU или using Metal в коде. Обратите внимание, что для CUDA загружать пакет cuDNN явным образом не требуется, но он должен быть установлен в среде.

Совместимость: Flux ≤ 0.13

В старых версиях Flux пакет CUDA.jl устанавливался автоматически для обеспечения поддержки GPU. Начиная с версии Flux 0.14 пакет CUDA.jl больше не является зависимостью и должен устанавливаться вручную.

Проверка доступности поддержки GPU

По умолчанию Flux автоматически проверяет систему на возможность поддержки функций GPU. Чтобы узнать, была ли ваша конфигурация GPU определена как допустимая, введите следующие команды:

julia> using CUDA

julia> CUDA.functional()
true

Для GPU AMD:

julia> using AMDGPU

julia> AMDGPU.functional()
true

julia> AMDGPU.functional(:MIOpen)
true

Для GPU Metal:

julia> using Metal

julia> Metal.functional()
true

Выбор бэкенда GPU

Допустимые бэкенды GPU: CUDA, AMDGPU и Metal.

Flux применяет Preferences.jl для выбора используемого по умолчанию бэкенда GPU.

Указать бэкенд можно двумя способами.

В REPL или коде вашего проекта вызовите Flux.gpu_backend!("AMDGPU") и при необходимости перезапустите сеанс Julia, чтобы изменения вступили в силу.
В файле LocalPreferences.toml в каталоге проекта укажите следующий параметр:

[Flux]
gpu_backend = "AMDGPU"

Текущий бэкенд GPU можно получить из переменной Flux.GPU_BACKEND:

julia> Flux.GPU_BACKEND
"CUDA"

Текущий бэкенд влияет на работу некоторых методов, таких как описанный ниже метод gpu.

Базовое использование GPU

Поддержка операций с массивами на другом оборудовании, таком как GPU, обеспечивается внешними пакетами, например CUDA.jl, AMDGPU.jl и Metal.jl. Работа Flux не зависит от типов массивов, поэтому достаточно перенести веса и данные модели в GPU.

Например, можно воспользоваться CUDA.CuArray (с преобразователем cu) для выполнения простейшего примера на GPU NVIDIA.

(Обратите внимание: для использования CUDA.CuArray должна быть доступна технология CUDA. Дополнительные сведения см. в инструкциях к CUDA.jl.)

using CUDA

W = cu(rand(2, 5)) # массив CuArray размером 2×5
b = cu(rand(2))

predict(x) = W*x .+ b
loss(x, y) = sum((predict(x) .- y).^2)

x, y = cu(rand(5)), cu(rand(2)) # Фиктивные данные
loss(x, y) # ~ 3

Обратите внимание, что в массив CUDA преобразуются как параметры (W, b), так и набор данных (x, y). Взятие производных и обучение выполняются как обычно.

Если определена структурированная модель, например слой Dense или Chain, достаточно преобразовать внутренние параметры. Flux предоставляет функцию fmap, которая позволяет изменять сразу все параметры модели.

d = Dense(10 => 5, σ)
d = fmap(cu, d)
d.weight # CuArray
d(cu(rand(10))) # Выходные данные CuArray

m = Chain(Dense(10 => 5, σ), Dense(5 => 2), softmax)
m = fmap(cu, m)
m(cu(rand(10)))

Для удобства Flux предоставляет функцию gpu, преобразующую модели и данные в формат GPU, если он поддерживается. По умолчанию она ничего не делает. Поэтому вы можете спокойно вызвать gpu для некоторых данных или модели (как показано ниже), и код не выдаст ошибку независимо от того, доступен ли GPU. Если библиотека GPU (например, CUDA) успешно загружена, функция gpu переместит данные с ЦП на GPU. Как показано ниже, при этом произойдут некоторые изменения типов, например, тип обычного массива изменится на CuArray.

julia> using Flux, CUDA

julia> m = Dense(10, 5) |> gpu
Dense(10 => 5)      # 55 параметров

julia> x = rand(10) |> gpu
10-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:
 0.066846445
 ⋮
 0.76706964

julia> m(x)
5-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:
 -0.99992573
 ⋮
 -0.547261

Доступна также аналогичная функция cpu для переноса моделей и данных с GPU обратно.

julia> x = rand(10) |> gpu
10-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:
 0.8019236
 ⋮
 0.7766742

julia> x |> cpu
10-element Vector{Float32}:
 0.8019236
 ⋮
 0.7766742

Передача обучающих данных

Для обучения модели с использованием GPU необходимо передать как саму модель, так и обучающие данные в память GPU. Переместить данные можно двумя разными способами.

Путем итерации по пакетам в объекте DataLoader с поочередной передачей каждого обучающего пакета в GPU. Этот способ рекомендуется для больших наборов данных. Вручную это может выглядеть так:
```
 train_loader = Flux.DataLoader((X, Y), batchsize=64, shuffle=true)
 # ... определение модели, настройка оптимизатора
 for epoch in 1:epochs
     for (x_cpu, y_cpu) in train_loader
         x = gpu(x_cpu)
         y = gpu(y_cpu)
         grads = gradient(m -> loss(m, x, y), model)
         Flux.update!(opt_state, model, grads[1])
     end
 end
```
Вместо того чтобы каждый раз писать этот код, можно просто вызвать gpu(::DataLoader):
```
 gpu_train_loader = Flux.DataLoader((X, Y), batchsize=64, shuffle=true) |> gpu
 # ... определение модели, настройка оптимизатора
 for epoch in 1:epochs
     for (x, y) in gpu_train_loader
         grads = gradient(m -> loss(m, x, y), model)
         Flux.update!(opt_state, model, grads[1])
     end
 end
```
Это эквивалентно DataLoader(MLUtils.mapobs(gpu, (X, Y)); keywords...). Нечто подобное можно также сделать с помощью CUDA.CuIterator, gpu_train_loader = CUDA.CuIterator(train_loader). Однако такой подход работает только с ограниченным набором типов данных: first(train_loader) должно быть кортежем (или NamedTuple) массивов.
Путем передачи сразу всех обучающих данных в GPU перед созданием DataLoader. Обычно такой подход применяется для наборов данных меньшего размера, которые гарантированно поместятся в доступную память GPU.
```
 gpu_train_loader = Flux.DataLoader((X, Y) |> gpu, batchsize = 32)
 # ...
 for epoch in 1:epochs
     for (x, y) in gpu_train_loader
         # ...
```
В данном случае (X, Y) |> gpu применяет функцию gpu к обоим массивам при рекурсии по структурам.

Сохранение обученных на GPU моделей

По завершении процесса обучения необходимо обязательно передать обученную модель обратно в область памяти cpu перед сериализацией или сохранением на диске. Как описывалось в предыдущем разделе, это можно сделать следующим образом:

model = cpu(model) # или model = model |> cpu

а затем выполнить такой код:

using BSON
# ...
BSON.@save "./path/to/trained_model.bson" model

# при таком подходе передаваемая в ЦП модель (на которую ссылается переменная `model`)
# существует только в пределах оператора `let`
let model = cpu(model)
   # ...
   BSON.@save "./path/to/trained_model.bson" model
end

# равносильно предыдущему варианту, но используется директива сохранения `key=value` из BSON.jl
BSON.@save "./path/to/trained_model.bson" model = cpu(model)

Причина в том, что модели, обученные на GPU, но не переданные в область памяти ЦП, ожидают в качестве входных данных массивов CuArray. Иными словами, модели Flux ожидают входных данных от устройства того же типа, на котором они были обучены.

В контролируемых сценариях, когда данные, подаваемые в загруженные модели, гарантированно размещаются в GPU, передавать их обратно в область памяти ЦП не требуется. Однако в рабочих средах, в которых артефакты используются совместно разными процессами, оборудованием или конфигурациями, нет гарантии того, что пакет CUDA.jl будет доступен процессу, выполняющему вывод на основе модели, загруженной с диска.

Отключение CUDA или выбор GPU, доступных Flux

Иногда бывает необходимо контролировать то, какие GPU доступны julia в системе с несколькими GPU, или полностью отключить GPU. Это можно делать с помощью переменной среды CUDA_VISIBLE_DEVICES.

Чтобы отключить все устройства:

$ export CUDA_VISIBLE_DEVICES='-1'

Чтобы выбрать определенные устройства по идентификаторам:

$ export CUDA_VISIBLE_DEVICES='0,1'

Дополнительные сведения об условном использовании GPU в CUDA.jl можно найти в документации, а использование переменной подробно описывается в записи блога об NVIDIA CUDA.

Использование объектов устройств

Для обеспечения более удобного синтаксиса Flux позволяет использовать объекты device GPU, с помощью которых можно легко передавать модели в GPU (если же бэкенд GPU недоступен, по умолчанию используется ЦП). Такой синтаксис имеет ряд преимуществ, включая автоматический выбор бэкенда GPU и стабильность типов при перемещении данных. С этой целью можно использовать функцию Flux.get_device.

Flux.get_device сначала проверяет предпочтительный GPU, а затем по возможности возвращает устройство для предпочтительного бэкенда. Возьмем следующий пример, в котором загружается пакет CUDA.jl для использования GPU NVIDIA ("CUDA" является предпочтительным бэкендом):

julia> using Flux, CUDA;

julia> device = Flux.get_device(; verbose=true)   # возвращает дескриптор GPU NVIDIA
[ Info: Using backend set in preferences: CUDA.
(::Flux.FluxCUDADevice) (generic function with 1 method)

julia> device.deviceID      # проверяем идентификатор GPU
CuDevice(0): NVIDIA GeForce GTX 1650

julia> model = Dense(2 => 3);

julia> model.weight     # модель изначально находится в памяти ЦП
3×2 Matrix{Float32}:
 -0.984794  -0.904345
  0.720379  -0.486398
  0.851011  -0.586942

julia> model = model |> device      # передаем модель в GPU
Dense(2 => 3)       # 9 параметров

julia> model.weight
3×2 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 -0.984794  -0.904345
  0.720379  -0.486398
  0.851011  -0.586942

Приоритетное устройство можно также задать с помощью функции Flux.gpu_backend!. Например, сначала зададим "CPU" в качестве приоритетного устройства:

julia> using Flux; Flux.gpu_backend!("CPU")
┌ Info: New GPU backend set: CPU.
└ Restart your Julia session for this change to take effect!

Затем после перезапуска сеанса Julia Flux.get_device возвращает дескриптор "CPU":

julia> using Flux, CUDA;    # даже если пакет CUDA загружен, все равно мы получаем устройство ЦП

julia> device = Flux.get_device(; verbose=true)   # получаем устройство ЦП
[ Info: Using backend set in preferences: CPU.
(::Flux.FluxCPUDevice) (generic function with 1 method)

julia> model = Dense(2 => 3);

julia> model = model |> device
Dense(2 => 3)       # 9 параметров

julia> model.weight     # ничего не изменилось; модель по-прежнему находится в ЦП
3×2 Matrix{Float32}:
 -0.942968   0.856258
  0.440009   0.714106
 -0.419192  -0.471838

Очевидно, это означает, что один и тот же код будет работать для любого бэкенда GPU, а также для ЦП.

Если приоритетный бэкенд недоступен или не работает, функция Flux.get_device ищет бэкенд CUDA, AMDGPU или Metal и возвращает соответствующее устройство (если бэкенд доступен и работает нормально). В противном случае возвращается устройство ЦП. В следующем примере предпочтительным GPU является "CUDA":

julia> using Flux;      # приоритетом является CUDA, но пакет CUDA.jl не загружен

julia> device = Flux.get_device(; verbose=true)       # в результате устройство выбирается автоматически
[ Info: Using backend set in preferences: CUDA.
┌ Warning: Trying to use backend: CUDA but it's trigger package is not loaded.
│ Please load the package and call this function again to respect the preferences backend.
└ @ Flux ~/fluxml/Flux.jl/src/functor.jl:637
[ Info: Using backend: CPU.
(::Flux.FluxCPUDevice) (generic function with 1 method)

Подробные сведения о выборе бэкенда см. в документации по Flux.get_device.

Перемещение данных между устройствами GPU

Flux также поддерживает получение дескрипторов определенных устройств GPU и передачу моделей с одного устройства GPU на другое, относящееся к тому же бэкенду. Попробуем сделать это с GPU NVIDIA. Сначала получаем список всех доступных устройств:

julia> using Flux, CUDA;

julia> CUDA.devices()
CUDA.DeviceIterator() for 3 devices:
0. GeForce RTX 2080 Ti
1. GeForce RTX 2080 Ti
2. TITAN X (Pascal)

Затем выбираем устройство с идентификатором 0:

julia> device0 = Flux.get_device("CUDA", 0)        # в настоящее время для бэкенда поддерживаются значения CUDA и AMDGPU
(::Flux.FluxCUDADevice) (generic function with 1 method)

Затем давайте перенесем простой плотный слой на GPU, представленный объектом device0:

julia> dense_model = Dense(2 => 3)
Dense(2 => 3)       # 9 параметров

julia> dense_model = dense_model |> device0;

julia> dense_model.weight
3×2 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
  0.695662   0.816299
 -0.204763  -0.10232
 -0.955829   0.538412

julia> CUDA.device(dense_model.weight)      # проверяем GPU, к которому прикреплена модель dense_model
CuDevice(0): GeForce RTX 2080 Ti

Затем получаем дескриптор устройства с идентификатором 1 и перемещаем dense_model на это устройство:

julia> device1 = Flux.get_device("CUDA", 1)
(::Flux.FluxCUDADevice) (generic function with 1 method)

julia> dense_model = dense_model |> device1;    # не выводите модель напрямую; см. предупреждение ниже

julia> CUDA.device(dense_model.weight)
CuDevice(1): GeForce RTX 2080 Ti

Из-за ограничения в Metal.jl в настоящее время такое перемещение данных между устройствами поддерживается только для бэкендов CUDA и AMDGPU.

Из-за ограничения в текущем механизме работы пакетов GPU вывод моделей в REPL после их перемещения на устройство GPU, отличное от текущего устройства, вызывает ошибку.

# Flux.AbstractDevice — Type

Flux.AbstractDevice <: Function

Абстрактный тип, представляющий объекты device для различных бэкендов GPU. В настоящее время поддерживаются бэкенды "CUDA", "AMDGPU", "Metal" и "CPU"; бэкенд "CPU" используется по умолчанию, если GPU недоступен. Расширения GPU во Flux определяют подтипы этого типа.

# Flux.FluxCPUDevice — Type

Flux.FluxCPUDevice <: Flux.AbstractDevice

Тип, представляющий объекты device для бэкенда "CPU" для Flux. Это вариант по умолчанию, если GPU недоступен для Flux.

# Flux.FluxCUDADevice — Type

FluxCUDADevice <: AbstractDevice

Тип, представляющий объекты device для бэкенда "CUDA" для Flux.

# Flux.FluxAMDGPUDevice — Type

FluxAMDGPUDevice <: AbstractDevice

Тип, представляющий объекты device для бэкенда "AMDGPU" для Flux.

# Flux.FluxMetalDevice — Type

FluxMetalDevice <: AbstractDevice

Тип, представляющий объекты device для бэкенда "Metal" для Flux.

# Flux.supported_devices — Function

Flux.supported_devices()

Возвращает все поддерживаемые бэкенды для Flux в порядке приоритета.

Пример

julia> using Flux;

julia> Flux.supported_devices()
("CUDA", "AMDGPU", "Metal", "CPU")

# Flux.get_device — Function

Flux.get_device(; verbose=false)::Flux.AbstractDevice

Возвращает объект device, соответствующий бэкенду, наиболее подходящему для текущего сеанса Julia.

Сначала эта функция проверяет, задан ли приоритетный бэкенд с помощью функции Flux.gpu_backend!. Если да, она пытается загрузить его. Если соответствующий активирующий пакет загружен и бэкенд функционален, загружается устройство device, соответствующее данному бэкенду. В противном случае бэкенд выбирается автоматически. Чтобы изменить приоритетный бэкенд, используйте функцию Flux.gpu_backend!.

Если приоритета нет, то эта функция проверяет по порядку каждый бэкенд ("CUDA", "AMDGPU", "Metal", "CPU") на предмет того, загружен ли он с помощью соответствующего активирующего пакета и функционален ли он. Если да, возвращается объект device, соответствующий бэкенду. Если не доступен ни один бэкенд GPU, возвращается Flux.FluxCPUDevice.

Если аргумент verbose имеет значение true, функция выводит информационные сообщения журнала.

Примеры

В приведенном ниже примере с помощью функции gpu_backend! в качестве приоритетного задан бэкенд "AMDGPU".

julia> using Flux;

julia> model = Dense(2 => 3)
Dense(2 => 3)       # 9 параметров

julia> device = Flux.get_device(; verbose=true)       # загружается устройство ЦП
[ Info: Using backend set in preferences: AMDGPU.
┌ Warning: Trying to use backend: AMDGPU but it's trigger package is not loaded.
│ Please load the package and call this function again to respect the preferences backend.
└ @ Flux ~/fluxml/Flux.jl/src/functor.jl:638
[ Info: Using backend: CPU.
(::Flux.FluxCPUDevice) (generic function with 1 method)

julia> model = model |> device
Dense(2 => 3)       # 9 параметров

julia> model.weight
3×2 Matrix{Float32}:
 -0.304362  -0.700477
 -0.861201   0.67825
 -0.176017   0.234188

Вот тот же пример, но с использованием "CUDA":

julia> using Flux, CUDA;

julia> model = Dense(2 => 3)
Dense(2 => 3)       # 9 параметров

julia> device = Flux.get_device(; verbose=true)
[ Info: Using backend set in preferences: AMDGPU.
┌ Warning: Trying to use backend: AMDGPU but it's trigger package is not loaded.
│ Please load the package and call this function again to respect the preferences backend.
└ @ Flux ~/fluxml/Flux.jl/src/functor.jl:637
[ Info: Using backend: CUDA.
(::Flux.FluxCUDADevice) (generic function with 1 method)

julia> model = model |> device
Dense(2 => 3)       # 9 параметров

julia> model.weight
3×2 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
  0.820013   0.527131
 -0.915589   0.549048
  0.290744  -0.0592499

Flux.get_device(backend::String, idx::Int = 0)::Flux.AbstractDevice

Возвращает объект устройства для бэкенда, заданного с помощью строки backend и idx. В настоящее время для backend поддерживаются значения "CUDA", "AMDGPU" и "CPU". Значением idx должно быть целое число от 0 до количества доступных устройств.

Примеры

julia> using Flux, CUDA;

julia> CUDA.devices()
CUDA.DeviceIterator() for 3 devices:
0. GeForce RTX 2080 Ti
1. GeForce RTX 2080 Ti
2. TITAN X (Pascal)

julia> device0 = Flux.get_device("CUDA", 0)
(::Flux.FluxCUDADevice) (generic function with 1 method)

julia> device0.deviceID
CuDevice(0): GeForce RTX 2080 Ti

julia> device1 = Flux.get_device("CUDA", 1)
(::Flux.FluxCUDADevice) (generic function with 1 method)

julia> device1.deviceID
CuDevice(1): GeForce RTX 2080 Ti

julia> cpu_device = Flux.get_device("CPU")
(::Flux.FluxCPUDevice) (generic function with 1 method)