Поддержка GPU

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

Для GPU AMD:

Для GPU Metal:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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